Downstream insertion tasks with NCF for tracking extrinsic contacts with tactile sensing
from dataclasses import dataclass, field
import yaml
import os
from typing import List


@dataclass
class Digit_VAE_Params:
    root_path: str = os.path.abspath(os.path.join(".."))
    image_size: int = 64
    channels: int = 3
    seq_len: int = 5
    latent_dim: int = 64
    enc_out_dim: int = 512
    source_data: str = "sim"
    checkpoint_dir = os.path.join(root_path, "digit_vae", "checkpoints")


@dataclass
class NCF_Params:
    root_path: str = os.path.abspath(os.path.join("."))
    seq_len: int = 5
    source_data: str = "sim"
    hidden_size_ncf: int = 128
    pc_subsample: float = 1.0
    vae_params = Digit_VAE_Params()
NCF_policies/algo2/ncf/config/config.py
carolinahiguera/NCF-Policy
import torch
import torch.nn as nn
import torch.nn.functional as F


class NCF(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.cfg = cfg
        self.seq_len = cfg.seq_len

        # digit encoder
        self.digit_emb_dim = cfg.vae_params.latent_dim
        digits_emb_in_dim = self.digit_emb_dim * self.seq_len
        digits_emb_out_dim = 128
        digit_emb_mid_dim = 128
        self.digits_down = nn.Sequential(
            nn.Linear(digits_emb_in_dim, digit_emb_mid_dim),
            nn.ReLU(),
            nn.Linear(digit_emb_mid_dim, digits_emb_out_dim),
            nn.ReLU(),
        )

        # ncf model
        ncf_dim = 2049
        input_dim = ncf_dim + (digits_emb_out_dim * 2) + (7 * 3)
        self.bn_in = nn.BatchNorm1d(input_dim)
        self.network = nn.Sequential(
            nn.Conv1d(input_dim, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Conv1d(512, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Conv1d(128, 1, 1),
            nn.Sigmoid(),
        )

    def forward(self, inputs):
        emb_ndf = inputs["emb_ndf"]
        images_left = inputs["digit_emb_left"]
        images_right = inputs["digit_emb_right"]
        ee_pose = inputs["ee_pose"][:, 0, [0, 1, 2, 4, 5, 6, 3]]
        ee_t1 = inputs["ee_pose_1"]
        ee_t2 = inputs["ee_pose_2"]
        contact_t1 = inputs["contact_oe_t1"].unsqueeze(dim=2)
        contact_t2 = inputs["contact_oe_t2"].unsqueeze(dim=2)

        images_left_down = self.digits_down(
            images_left.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        images_right_down = self.digits_down(
            images_right.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        digit_feat_all = torch.cat([images_left_down, images_right_down], dim=1)

        n_pts = emb_ndf.shape[1]
        ee_t0 = ee_pose[:, None, :].repeat(1, n_pts, 1)
        ee_t1 = ee_t1[:, None, :].repeat(1, n_pts, 1)
        ee_t2 = ee_t2[:, None, :].repeat(1, n_pts, 1)
        digit_feat_all = digit_feat_all[:, None, :].repeat(1, n_pts, 1)
        # x_in = torch.cat((emb_ndf, contact_t1, contact_t2, ee_t1, ee_t2), 2)
        x_in = torch.cat((emb_ndf, digit_feat_all, ee_t0, ee_t1, ee_t2), 2)
        x_in = x_in.permute(0, 2, 1)

        x_in = self.bn_in(x_in)
        out = self.network(x_in)
        out = out.squeeze(dim=1)
        return out

NCF_policies/algo2/ncf/contact_regressor/ncf_conv.py
carolinahiguera/NCF-Policy
import torch
import torch.nn as nn
import torch.autograd as autograd
import torch.nn.functional as F


class NCF(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.cfg = cfg
        self.seq_len = cfg.seq_len

        # digit encoder
        self.digit_emb_dim = cfg.vae_params.latent_dim
        digits_emb_in_dim = self.digit_emb_dim * self.seq_len
        digits_emb_out_dim = 128
        digit_emb_mid_dim = 128
        self.digits_down = nn.Sequential(
            nn.Linear(digits_emb_in_dim, digit_emb_mid_dim),
            nn.ReLU(),
            nn.Linear(digit_emb_mid_dim, digits_emb_out_dim),
            nn.ReLU(),
        )

        # ncf model
        ncf_dim = 768 + 3
        input_dim = ncf_dim + (digits_emb_out_dim * 2) + (7 * 3)
        self.bn_in = nn.BatchNorm1d(input_dim)
        # self.network = nn.Sequential(
        #     nn.Conv1d(input_dim, 512, 1),
        #     nn.BatchNorm1d(512),
        #     nn.ReLU(),
        #     nn.Conv1d(512, 128, 1),
        #     nn.BatchNorm1d(128),
        #     nn.ReLU(),
        #     nn.Conv1d(128, 1, 1),
        #     nn.Sigmoid(),
        # )
        self.network = nn.Sequential(
            nn.Linear(input_dim, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Linear(128, 1, 1),
            nn.Sigmoid(),
        )

    def forward(self, inputs):
        images_left = inputs["digit_emb_left"]
        images_right = inputs["digit_emb_right"]
        # ee_pose = inputs["ee_pose"][:, 0, [0, 1, 2, 4, 5, 6, 3]]
        ee_pose = inputs["ee_pose"][:, 0]
        ee_t1 = inputs["ee_pose_1"]
        ee_t2 = inputs["ee_pose_2"]
        shape_emb = inputs["shape_emb"]
        ncf_query_points = inputs["ncf_query_points"]

        images_left_down = self.digits_down(
            images_left.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        images_right_down = self.digits_down(
            images_right.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        digit_feat_all = torch.cat([images_left_down, images_right_down], dim=1)

        n_pts = ncf_query_points.shape[1]
        batch_size = ncf_query_points.shape[0]
        ee_t0 = ee_pose[:, None, :].repeat(1, n_pts, 1)
        ee_t1 = ee_t1[:, None, :].repeat(1, n_pts, 1)
        ee_t2 = ee_t2[:, None, :].repeat(1, n_pts, 1)
        digit_feat_all = digit_feat_all[:, None, :].repeat(1, n_pts, 1)
        shape_emb = shape_emb[:, None, :].repeat(1, n_pts, 1)
        # x_in = torch.cat((emb_ndf, contact_t1, contact_t2, ee_t1, ee_t2), 2)
        x_in = torch.cat(
            (ncf_query_points, shape_emb, digit_feat_all, ee_t0, ee_t1, ee_t2), 2
        )
        x_in = x_in.permute(0, 2, 1)

        x_in = self.bn_in(x_in).permute(0, 2, 1)
        out = self.network(x_in.flatten(0, 1))
        out = out.unflatten(0, (batch_size, n_pts))
        out = out.squeeze(dim=2)
        return out

NCF_policies/algo2/ncf/contact_regressor/ncf_mlp.py
carolinahiguera/NCF-Policy
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class TransformerConv(nn.Module):
    def __init__(self, input_dim, num_heads, hidden_dim, num_layers):
        super(TransformerConv, self).__init__()

        self.input_dim = input_dim
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        # self.positional_encoding = PositionalEncoding(input_dim)

        self.conv_layer1 = TransformerConvAttentionLayer(
            input_dim, hidden_dim[0], num_heads
        )
        # self.conv_layer2 = TransformerConvAttentionLayer(
        #     hidden_dim[0], hidden_dim[1], num_heads
        # )
        self.reg = nn.Sequential(
            nn.Conv1d(hidden_dim[0], hidden_dim[1], 1),
            nn.BatchNorm1d(hidden_dim[1]),
            nn.ReLU(inplace=True),
            nn.Conv1d(hidden_dim[1], 1, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        # x = self.positional_encoding(x)
        x = x.permute(0, 2, 1)  # Reshape to (batch, d, n)
        x = self.conv_layer1(x)
        # x = self.conv_layer2(x)
        x = self.reg(x)
        x = x.squeeze(1)
        return x


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[: x.size(0), :]
        return x


class TransformerConvAttentionLayer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_heads, dropout=0.2):
        super(TransformerConvAttentionLayer, self).__init__()

        self.conv1 = nn.Conv1d(input_dim, hidden_dim, 1)
        # self.conv2 = nn.Conv1d(hidden_dim, input_dim, 1)
        self.norm1 = nn.BatchNorm1d(input_dim)
        self.norm2 = nn.BatchNorm1d(hidden_dim)
        self.dropout = nn.Dropout(dropout)
        # self.attention = nn.MultiheadAttention(hidden_dim, num_heads)

        self.linear_q = nn.Conv1d(input_dim, input_dim, 1)
        self.linear_k = nn.Conv1d(input_dim, input_dim, 1)
        self.linear_v = nn.Conv1d(input_dim, input_dim, 1)

    def scaled_dot_product(self, q, k, v, mask=None):
        q = q.permute(0, 2, 1)
        k = k.permute(0, 2, 1)
        v = v.permute(0, 2, 1)
        d_k = q.size()[-1]
        attn_logits = torch.matmul(q, k.transpose(-2, -1))
        attn_logits = attn_logits / math.sqrt(d_k)
        if mask is not None:
            attn_logits = attn_logits.masked_fill(mask == 0, -9e15)
        attention = F.softmax(attn_logits, dim=-1)
        values = torch.matmul(attention, v)
        values = values.permute(0, 2, 1)
        attention = attention.permute(0, 2, 1)
        return values, attention

    def forward(self, x):
        x_q, x_k, x_v = self.linear_q(x), self.linear_k(x), self.linear_v(x)
        attn_output, _ = self.scaled_dot_product(x_q, x_k, x_v)

        x = x + self.dropout(attn_output)
        x = self.norm1(x)

        x = self.conv1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.norm2(x)

        return x


class NCF(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.cfg = cfg
        self.seq_len = cfg.seq_len

        # digit encoder
        self.digit_emb_dim = cfg.vae_params.latent_dim
        digits_emb_in_dim = self.digit_emb_dim * self.seq_len
        digits_emb_out_dim = 128
        digit_emb_mid_dim = 128
        self.digits_down = nn.Sequential(
            nn.Linear(digits_emb_in_dim, digit_emb_mid_dim),
            nn.ReLU(),
            nn.Linear(digit_emb_mid_dim, digits_emb_out_dim),
            nn.ReLU(),
        )

        # ncf model
        ncf_dim = 768 + 3
        input_dim = ncf_dim + (digits_emb_out_dim * 2) + (7 * 3)
        hidden_dim = [512, 128]
        self.bn_in = nn.BatchNorm1d(input_dim)
        self.network = TransformerConv(
            input_dim=input_dim, num_heads=2, hidden_dim=hidden_dim, num_layers=2
        )

    def forward(self, inputs):
        images_left = inputs["digit_emb_left"]
        images_right = inputs["digit_emb_right"]
        # ee_pose = inputs["ee_pose"][:, 0, [0, 1, 2, 4, 5, 6, 3]]
        ee_pose = inputs["ee_pose"][:, 0]
        ee_t1 = inputs["ee_pose_1"]
        ee_t2 = inputs["ee_pose_2"]
        shape_emb = inputs["shape_emb"]
        ncf_query_points = inputs["ncf_query_points"]

        images_left_down = self.digits_down(
            images_left.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        images_right_down = self.digits_down(
            images_right.reshape(-1, self.seq_len * self.digit_emb_dim)
        )
        digit_feat_all = torch.cat([images_left_down, images_right_down], dim=1)

        n_pts = ncf_query_points.shape[1]
        ee_t0 = ee_pose[:, None, :].repeat(1, n_pts, 1)
        ee_t1 = ee_t1[:, None, :].repeat(1, n_pts, 1)
        ee_t2 = ee_t2[:, None, :].repeat(1, n_pts, 1)
        digit_feat_all = digit_feat_all[:, None, :].repeat(1, n_pts, 1)

        shape_emb = shape_emb[:, None, :].repeat(1, n_pts, 1)
        # x_in = torch.cat((emb_ndf, contact_t1, contact_t2, ee_t1, ee_t2), 2)
        # x_in = torch.cat((emb_ndf, digit_feat_all, ee_t0, ee_t1, ee_t2), 2)
        x_in = torch.cat(
            (ncf_query_points, shape_emb, digit_feat_all, ee_t0, ee_t1, ee_t2), 2
        )

        out = self.network(x_in)

        return out

NCF_policies/algo2/ncf/contact_regressor/ncf_transformer.py
carolinahiguera/NCF-Policy
from typing import Any
import pytorch_lightning as pl

import torch
import torch.nn as nn
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import cv2
import numpy as np

# import trimesh
# import vedo
# from vedo import Points, show, trimesh2vedo
from scipy.spatial.transform import Rotation as R

DEBUG = False


def show_obj_probabilities(pc, mesh_file, dist, interactive=False):
    object_3d_cam = dict(
        position=(-7.54483e-3, -0.0849045, -0.250212),
        focal_point=(-4.82255e-3, -2.87705e-3, 0),
        viewup=(0.580505, -0.775588, 0.247946),
        distance=0.263329,
        clipping_range=(0.174192, 0.376170),
    )

    pc = Points(pc, r=15)
    pc = pc.cmap(
        "plasma",
        dist,
        vmin=0.0,
        vmax=1.0,
    )

    object_trimesh = trimesh.load_mesh(mesh_file)
    T = np.eye(4)
    T[0:3, 0:3] = R.from_euler("xyz", [0, 0, 90], degrees=True).as_matrix()
    object_trimesh = object_trimesh.apply_transform(T)
    mesh_vedo = trimesh2vedo(object_trimesh).clone()
    # mesh_vedo.subdivide(n=10, method=2)
    mesh_vedo = mesh_vedo.interpolate_data_from(
        pc, n=5, on="points", kernel="gaussian"
    ).cmap("plasma", vmin=0.0, vmax=1.0)

    if interactive:
        show([mesh_vedo, pc], axes=1, camera=object_3d_cam)
    else:
        plt = vedo.Plotter(offscreen=True, size=(500, 500))
        plt.show([mesh_vedo, pc]).screenshot("contact.png").close()
        print("")
        # img = show(
        #     [mesh_vedo, pc],
        #     axes=1,
        #     camera=object_3d_cam,
        #     interactive=False,
        #     offscreen=True,
        # ).screenshot("contact.png")
        # ).screenshot(asarray=True)
        # return img


class NCF_Pipeline(pl.LightningModule):
    def __init__(self, cfg_ncf, cfg_train, digit_vae, ncf, ndf):
        super().__init__()

        self.cfg_ncf = cfg_ncf
        self.cfg_train = cfg_train
        self.seq_len = cfg_ncf.seq_len
        self.digit_emb_dim = cfg_ncf.vae_params.latent_dim

        self.digit_vae = digit_vae
        self.ncf = ncf
        self.ndf = ndf

        self.mse_loss = nn.MSELoss(reduction="none")
        self.bce_loss = nn.BCELoss(reduction="none")
        self.l1_loss = nn.L1Loss(reduction="none")

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.cfg_train.lr)

    # def configure_optimizers(self):
    #     # Using a scheduler is optional but can be helpful.
    #     # The scheduler reduces the LR if the validation performance hasn't improved for the last N epochs
    #     # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    #     #     optimizer, mode="min", factor=0.2, patience=20, min_lr=5e-5
    #     # )

    #     lr = (
    #         self.cfg_train.lr
    #         if self.cfg_train.ncf_arch == "transformer"
    #         else self.cfg_train.lr * 0.5
    #     )

    #     optimizer = torch.optim.Adam(self.parameters(), lr=lr)

    #     if self.cfg_train.ncf_arch == "transformer":
    #         scheduler = torch.optim.lr_scheduler.StepLR(
    #             optimizer, step_size=5, gamma=0.1
    #         )
    #     else:
    #         scheduler = torch.optim.lr_scheduler.StepLR(
    #             optimizer, step_size=2, gamma=0.1
    #         )

    #     return {
    #         "optimizer": optimizer,
    #         "lr_scheduler": {
    #             "scheduler": scheduler,
    #         },
    #     }
    #     # return {"optimizer": optimizer, "monitor": "val_loss"}

    def _get_digit_embeddings(self, images):
        batch_size = images.shape[0]
        seq_len = images.shape[1]
        images_hat = torch.zeros_like(images)

        digit_embeddings = torch.zeros(self.seq_len, batch_size, self.digit_emb_dim).to(
            images.device
        )

        with torch.no_grad():
            for i in range(seq_len):
                s = images[:, i, :, :, :].unsqueeze(dim=1)
                self.digit_vae.eval()
                x_hat, emb = self.digit_vae(s)
                images_hat[:, i, :, :, :] = x_hat
                digit_embeddings[i, :, :] = emb

        return digit_embeddings.permute(1, 0, 2), images_hat

    def compute_loss_mse(self, gt, pred):
        loss = self.mse_loss(pred, gt)
        return loss.mean()

    def compute_loss_l1(self, gt, pred):
        loss = self.l1_loss(pred, gt)
        return loss.mean()

    def compute_loss_bce(self, gt, pred):
        pred = pred.squeeze()
        gt = gt.squeeze()
        loss = self.bce_loss(pred, gt)
        return loss.mean()

    def compute_loss_classic(self, gt, pred):
        label = gt.squeeze()
        # label = (label + 1) / 2.
        loss = (
            -1
            * (
                label * torch.log(pred + 1e-5)
                + (1 - label) * torch.log(1 - pred + 1e-5)
            ).mean()
        )
        return loss

    def compute_loss(self, pred, batch):
        gt = batch["p_contact_t0"]

        if self.cfg_train.loss == "mse":
            loss = self.compute_loss_mse(gt, pred)
        elif self.cfg_train.loss == "l1":
            loss = self.compute_loss_l1(gt, pred)
        elif self.cfg_train.loss == "bce":
            loss = self.compute_loss_bce(gt, pred)
        elif self.cfg_train.loss == "classic":
            loss = self.compute_loss_classic(gt, pred)
        else:
            return NotImplementedError
        return loss

    def forward(self, batch):
        images_left = batch["digit_imgs_left"]
        images_right = batch["digit_imgs_right"]
        ndf_point_cloud = batch["ndf_point_cloud"]
        ncf_query_points = batch["ncf_query_points"]

        # get digit embeddings
        embs_left, images_hat_left = self._get_digit_embeddings(images_left)
        embs_right, images_hat_right = self._get_digit_embeddings(images_right)

        with torch.no_grad():
            shape_emb = self.ndf(ndf_point_cloud)

        inputs_ncf = {
            "digit_emb_left": embs_left,
            "digit_emb_right": embs_right,
            "ee_pose": batch["ee_pose"],
            "ee_pose_1": batch["ee_pose_1"],
            "ee_pose_2": batch["ee_pose_2"],
            "shape_emb": shape_emb,
            "ncf_query_points": ncf_query_points,
        }

        pred_contact = self.ncf(inputs_ncf)

        out = {
            "pred_contact": pred_contact,
            "images_hat_left": images_hat_left,
            "images_hat_right": images_hat_right,
        }
        return out

        # return pred_contact

    def plot_contact(self, pred_contact, batch, interactive):
        idx = 0
        gt = batch["p_contact_t0"][idx].cpu().numpy()
        idx_query = batch["idx_query"][idx].cpu().numpy()
        pred = pred_contact[idx].detach().cpu().numpy()
        pc = np.load(batch["pc_file"][idx])
        mesh_file = batch["mesh_file"][idx]

        gt_full = np.zeros((pc.shape[0]))
        gt_full[idx_query] = gt
        pred_full = np.zeros((pc.shape[0]))
        pred_full[idx_query] = pred

        gt_img = show_obj_probabilities(pc, mesh_file, gt_full, interactive)

        if not interactive:
            pred_img = show_obj_probabilities(pc, pred_full)
            img = cv2.hconcat([gt_img, pred_img])
            plt.imshow(img)
            plt.show()

    def training_step(self, batch, batch_idx):
        out_ncf = self.forward(batch)
        pred_contact_oe = out_ncf["pred_contact"]
        loss = self.compute_loss(pred_contact_oe, batch)
        batch_size = pred_contact_oe.shape[0]
        # print(loss.item())
        if DEBUG:
            self.plot_contact(pred_contact_oe, batch, interactive=False)
        self.log(
            "train_loss",
            loss,
            on_step=True,
            on_epoch=True,
            prog_bar=False,
            logger=True,
            batch_size=batch_size,
        )
        return {"loss": loss, "pred_contact": pred_contact_oe}

    def validation_step(self, batch, batch_idx):
        out_ncf = self.forward(batch)
        pred_contact_oe = out_ncf["pred_contact"]
        loss = self.compute_loss(pred_contact_oe, batch)
        batch_size = pred_contact_oe.shape[0]
        self.log(
            "val_loss",
            loss,
            on_step=True,
            on_epoch=True,
            prog_bar=False,
            logger=True,
            batch_size=batch_size,
        )
        return {"loss": loss, "pred_contact": pred_contact_oe}

NCF_policies/algo2/ncf/contact_regressor/pipeline.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation
# https://arxiv.org/abs/2112.05124
# Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B. Tenenbaum, Alberto Rodriguez, 
# Pulkit Agrawal, Vincent Sitzmann
# --------------------------------------------------------

import os
import sys
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

EPS = 1e-6
USE_CUDA = True
# device = torch.device("cuda:1") if USE_CUDA else torch.device("cpu")
device = torch.device("cuda:0") if USE_CUDA else torch.device("cpu")


# print("=============================")
# print(f'layers equi device {device}')
# print("=============================")



def knn(x, k):
    inner = -2 * torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    idx = pairwise_distance.topk(k=k, dim=-1)[1]  # (batch_size, num_points, k)
    return idx


def get_graph_feature(x, k=20, idx=None, x_coord=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        if x_coord is not None:  # dynamic knn graph
            idx = knn(x_coord, k=k)  # (batch_size, num_points, k)
        else:  # fixed knn graph with input point coordinates
            idx = knn(x, k=k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3)
    x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)

    feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 4, 1, 2).contiguous()

    return feature


def get_graph_feature_cross(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        idx = knn(x, k=k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3)
    x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
    cross = torch.cross(feature, x, dim=-1)

    feature = (
        torch.cat((feature - x, x, cross), dim=3).permute(0, 3, 4, 1, 2).contiguous()
    )

    return feature


def get_graph_mean(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.reshape(batch_size, -1, num_points).contiguous()
    if idx is None:
        idx = knn(x, k=k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3).mean(
        2, keepdim=False
    )
    x = x.view(batch_size, num_points, num_dims, 3)

    feature = (feature - x).permute(0, 2, 3, 1).contiguous()

    return feature


def get_shell_mean_cross(x, k=10, nk=4, idx_all=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.reshape(batch_size, -1, num_points).contiguous()
    if idx_all is None:
        idx_all = knn(x, k=nk * k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = []
    for i in range(nk):
        idx.append(idx_all[:, :, i * k : (i + 1) * k])
        idx[i] = idx[i] + idx_base
        idx[i] = idx[i].view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    x = x.view(batch_size, num_points, num_dims, 3)
    feature = []
    for i in range(nk):
        feature.append(x.view(batch_size * num_points, -1)[idx[i], :])
        feature[i] = (
            feature[i]
            .view(batch_size, num_points, k, num_dims, 3)
            .mean(2, keepdim=False)
        )
        feature[i] = feature[i] - x
        cross = torch.cross(feature[i], x, dim=3)
        feature[i] = torch.cat((feature[i], cross), dim=2)

    feature = torch.cat(feature, dim=2).permute(0, 2, 3, 1).contiguous()

    return feature


class VNLinear(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(VNLinear, self).__init__()
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        x_out = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        return x_out


class VNLeakyReLU(nn.Module):
    def __init__(self, in_channels, share_nonlinearity=False, negative_slope=0.2):
        super(VNLeakyReLU, self).__init__()
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (x * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * x + (1 - self.negative_slope) * (
            mask * x + (1 - mask) * (x - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNLinearLeakyReLU(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        dim=5,
        share_nonlinearity=False,
        use_batchnorm=True,
        negative_slope=0.2,
    ):
        super(VNLinearLeakyReLU, self).__init__()
        self.dim = dim
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm
        self.negative_slope = negative_slope

        # Conv
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

        # BatchNorm
        self.use_batchnorm = use_batchnorm
        if use_batchnorm == True:
            self.batchnorm = VNBatchNorm(out_channels, dim=dim)

        # LeakyReLU
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, out_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        # Conv
        p = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        # InstanceNorm
        if self.use_batchnorm == True:
            p = self.batchnorm(p)
        # LeakyReLU
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (p * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * p + (1 - self.negative_slope) * (
            mask * p + (1 - mask) * (p - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNBatchNorm(nn.Module):
    def __init__(self, num_features, dim):
        super(VNBatchNorm, self).__init__()
        self.dim = dim
        if dim == 3 or dim == 4:
            self.bn = nn.BatchNorm1d(num_features)
        elif dim == 5:
            self.bn = nn.BatchNorm2d(num_features)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        norm = torch.sqrt((x * x).sum(2))
        norm_bn = self.bn(norm)
        norm = norm.unsqueeze(2)
        norm_bn = norm_bn.unsqueeze(2)
        x = x / norm * norm_bn

        return x


class VNMaxPool(nn.Module):
    def __init__(self, in_channels, share_nonlinearity=False):
        super(VNMaxPool, self).__init__()
        if share_nonlinearity:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (x * d).sum(2, keepdim=True)
        idx = dotprod.max(dim=-1, keepdim=False)[1]
        index_tuple = torch.meshgrid([torch.arange(j) for j in x.size()[:-1]]) + (idx,)
        x_max = x[index_tuple]
        return x_max


class VNStdFeature(nn.Module):
    def __init__(
        self,
        in_channels,
        dim=4,
        normalize_frame=False,
        share_nonlinearity=False,
        use_batchnorm=True,
    ):
        super(VNStdFeature, self).__init__()
        self.dim = dim
        self.normalize_frame = normalize_frame
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm

        self.vn1 = VNLinearLeakyReLU(
            in_channels,
            in_channels // 2,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        self.vn2 = VNLinearLeakyReLU(
            in_channels // 2,
            in_channels // 4,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        if normalize_frame:
            self.vn_lin = nn.Linear(in_channels // 4, 2, bias=False)
        else:
            self.vn_lin = nn.Linear(in_channels // 4, 3, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        z0 = self.vn1(x)
        z0 = self.vn2(z0)
        z0 = self.vn_lin(z0.transpose(1, -1)).transpose(1, -1)

        if self.normalize_frame:
            # make z0 orthogonal. u2 = v2 - proj_u1(v2)
            v1 = z0[:, 0, :]
            # u1 = F.normalize(v1, dim=1)
            v1_norm = torch.sqrt((v1 * v1).sum(1, keepdim=True))
            u1 = v1 / (v1_norm + EPS)
            v2 = z0[:, 1, :]
            v2 = v2 - (v2 * u1).sum(1, keepdim=True) * u1
            # u2 = F.normalize(u2, dim=1)
            v2_norm = torch.sqrt((v2 * v2).sum(1, keepdim=True))
            u2 = v2 / (v2_norm + EPS)

            # compute the cross product of the two output vectors
            u3 = torch.cross(u1, u2)
            z0 = torch.stack([u1, u2, u3], dim=1).transpose(1, 2)
        else:
            z0 = z0.transpose(1, 2)

        if self.dim == 4:
            x_std = torch.einsum("bijm,bjkm->bikm", x, z0)
        elif self.dim == 3:
            x_std = torch.einsum("bij,bjk->bik", x, z0)
        elif self.dim == 5:
            x_std = torch.einsum("bijmn,bjkmn->bikmn", x, z0)

        return x_std, z0


# Resnet Blocks
class VNResnetBlockFC(nn.Module):
    """Fully connected ResNet Block class.

    Args:
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
    """

    def __init__(self, size_in, size_out=None, size_h=None):
        super().__init__()
        # Attributes
        if size_out is None:
            size_out = size_in

        if size_h is None:
            size_h = min(size_in, size_out)

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        self.fc_0 = VNLinear(size_in, size_h)
        self.fc_1 = VNLinear(size_h, size_out)
        self.actvn_0 = VNLeakyReLU(
            size_in, negative_slope=0.2, share_nonlinearity=False
        )
        self.actvn_1 = VNLeakyReLU(size_h, negative_slope=0.2, share_nonlinearity=False)

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = VNLinear(size_in, size_out)
        # Initialization
        nn.init.zeros_(self.fc_1.map_to_feat.weight)

    def forward(self, x):
        net = self.fc_0(self.actvn_0(x))
        dx = self.fc_1(self.actvn_1(net))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx

NCF_policies/algo2/ncf/ndf/layers_equi.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation
# https://arxiv.org/abs/2112.05124
# Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B. Tenenbaum, Alberto Rodriguez, 
# Pulkit Agrawal, Vincent Sitzmann
# --------------------------------------------------------

import os
import sys
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

EPS = 1e-6
USE_CUDA = True
device = torch.device("cuda:1") if USE_CUDA else torch.device("cpu")


# print("=============================")
# print(f'layers equi device {device}')
# print("=============================")



def knn(x, k):
    inner = -2 * torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    idx = pairwise_distance.topk(k=k, dim=-1)[1]  # (batch_size, num_points, k)
    return idx


def get_graph_feature(x, k=20, idx=None, x_coord=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        if x_coord is not None:  # dynamic knn graph
            idx = knn(x_coord, k=k)  # (batch_size, num_points, k)
        else:  # fixed knn graph with input point coordinates
            idx = knn(x, k=k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3)
    x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)

    feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 4, 1, 2).contiguous()

    return feature


def get_graph_feature_cross(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        idx = knn(x, k=k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3)
    x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
    cross = torch.cross(feature, x, dim=-1)

    feature = (
        torch.cat((feature - x, x, cross), dim=3).permute(0, 3, 4, 1, 2).contiguous()
    )

    return feature


def get_graph_mean(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.reshape(batch_size, -1, num_points).contiguous()
    if idx is None:
        idx = knn(x, k=k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = idx + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    feature = x.view(batch_size * num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims, 3).mean(
        2, keepdim=False
    )
    x = x.view(batch_size, num_points, num_dims, 3)

    feature = (feature - x).permute(0, 2, 3, 1).contiguous()

    return feature


def get_shell_mean_cross(x, k=10, nk=4, idx_all=None):
    batch_size = x.size(0)
    num_points = x.size(3)
    x = x.reshape(batch_size, -1, num_points).contiguous()
    if idx_all is None:
        idx_all = knn(x, k=nk * k)  # (batch_size, num_points, k)
    # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points

    idx = []
    for i in range(nk):
        idx.append(idx_all[:, :, i * k : (i + 1) * k])
        idx[i] = idx[i] + idx_base
        idx[i] = idx[i].view(-1)

    _, num_dims, _ = x.size()
    num_dims = num_dims // 3

    x = x.transpose(
        2, 1
    ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
    x = x.view(batch_size, num_points, num_dims, 3)
    feature = []
    for i in range(nk):
        feature.append(x.view(batch_size * num_points, -1)[idx[i], :])
        feature[i] = (
            feature[i]
            .view(batch_size, num_points, k, num_dims, 3)
            .mean(2, keepdim=False)
        )
        feature[i] = feature[i] - x
        cross = torch.cross(feature[i], x, dim=3)
        feature[i] = torch.cat((feature[i], cross), dim=2)

    feature = torch.cat(feature, dim=2).permute(0, 2, 3, 1).contiguous()

    return feature


class VNLinear(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(VNLinear, self).__init__()
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        x_out = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        return x_out


class VNLeakyReLU(nn.Module):
    def __init__(self, in_channels, share_nonlinearity=False, negative_slope=0.2):
        super(VNLeakyReLU, self).__init__()
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (x * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * x + (1 - self.negative_slope) * (
            mask * x + (1 - mask) * (x - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNLinearLeakyReLU(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        dim=5,
        share_nonlinearity=False,
        use_batchnorm=True,
        negative_slope=0.2,
    ):
        super(VNLinearLeakyReLU, self).__init__()
        self.dim = dim
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm
        self.negative_slope = negative_slope

        # Conv
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

        # BatchNorm
        self.use_batchnorm = use_batchnorm
        if use_batchnorm == True:
            self.batchnorm = VNBatchNorm(out_channels, dim=dim)

        # LeakyReLU
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, out_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        # Conv
        p = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        # InstanceNorm
        if self.use_batchnorm == True:
            p = self.batchnorm(p)
        # LeakyReLU
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (p * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * p + (1 - self.negative_slope) * (
            mask * p + (1 - mask) * (p - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNBatchNorm(nn.Module):
    def __init__(self, num_features, dim):
        super(VNBatchNorm, self).__init__()
        self.dim = dim
        if dim == 3 or dim == 4:
            self.bn = nn.BatchNorm1d(num_features)
        elif dim == 5:
            self.bn = nn.BatchNorm2d(num_features)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        norm = torch.sqrt((x * x).sum(2))
        norm_bn = self.bn(norm)
        norm = norm.unsqueeze(2)
        norm_bn = norm_bn.unsqueeze(2)
        x = x / norm * norm_bn

        return x


class VNMaxPool(nn.Module):
    def __init__(self, in_channels, share_nonlinearity=False):
        super(VNMaxPool, self).__init__()
        if share_nonlinearity:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (x * d).sum(2, keepdim=True)
        idx = dotprod.max(dim=-1, keepdim=False)[1]
        index_tuple = torch.meshgrid([torch.arange(j) for j in x.size()[:-1]]) + (idx,)
        x_max = x[index_tuple]
        return x_max


class VNStdFeature(nn.Module):
    def __init__(
        self,
        in_channels,
        dim=4,
        normalize_frame=False,
        share_nonlinearity=False,
        use_batchnorm=True,
    ):
        super(VNStdFeature, self).__init__()
        self.dim = dim
        self.normalize_frame = normalize_frame
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm

        self.vn1 = VNLinearLeakyReLU(
            in_channels,
            in_channels // 2,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        self.vn2 = VNLinearLeakyReLU(
            in_channels // 2,
            in_channels // 4,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        if normalize_frame:
            self.vn_lin = nn.Linear(in_channels // 4, 2, bias=False)
        else:
            self.vn_lin = nn.Linear(in_channels // 4, 3, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        z0 = self.vn1(x)
        z0 = self.vn2(z0)
        z0 = self.vn_lin(z0.transpose(1, -1)).transpose(1, -1)

        if self.normalize_frame:
            # make z0 orthogonal. u2 = v2 - proj_u1(v2)
            v1 = z0[:, 0, :]
            # u1 = F.normalize(v1, dim=1)
            v1_norm = torch.sqrt((v1 * v1).sum(1, keepdim=True))
            u1 = v1 / (v1_norm + EPS)
            v2 = z0[:, 1, :]
            v2 = v2 - (v2 * u1).sum(1, keepdim=True) * u1
            # u2 = F.normalize(u2, dim=1)
            v2_norm = torch.sqrt((v2 * v2).sum(1, keepdim=True))
            u2 = v2 / (v2_norm + EPS)

            # compute the cross product of the two output vectors
            u3 = torch.cross(u1, u2)
            z0 = torch.stack([u1, u2, u3], dim=1).transpose(1, 2)
        else:
            z0 = z0.transpose(1, 2)

        if self.dim == 4:
            x_std = torch.einsum("bijm,bjkm->bikm", x, z0)
        elif self.dim == 3:
            x_std = torch.einsum("bij,bjk->bik", x, z0)
        elif self.dim == 5:
            x_std = torch.einsum("bijmn,bjkmn->bikmn", x, z0)

        return x_std, z0


# Resnet Blocks
class VNResnetBlockFC(nn.Module):
    """Fully connected ResNet Block class.

    Args:
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
    """

    def __init__(self, size_in, size_out=None, size_h=None):
        super().__init__()
        # Attributes
        if size_out is None:
            size_out = size_in

        if size_h is None:
            size_h = min(size_in, size_out)

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        self.fc_0 = VNLinear(size_in, size_h)
        self.fc_1 = VNLinear(size_h, size_out)
        self.actvn_0 = VNLeakyReLU(
            size_in, negative_slope=0.2, share_nonlinearity=False
        )
        self.actvn_1 = VNLeakyReLU(size_h, negative_slope=0.2, share_nonlinearity=False)

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = VNLinear(size_in, size_out)
        # Initialization
        nn.init.zeros_(self.fc_1.map_to_feat.weight)

    def forward(self, x):
        net = self.fc_0(self.actvn_0(x))
        dx = self.fc_1(self.actvn_1(net))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx

NCF_policies/algo2/ncf/ndf/layers_equi2.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation
# https://arxiv.org/abs/2112.05124
# Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B. Tenenbaum, Alberto Rodriguez, 
# Pulkit Agrawal, Vincent Sitzmann
# --------------------------------------------------------

import torch
import torch.nn as nn

EPS = 1e-6


def meanpool(x, dim=-1, keepdim=False):
    out = x.mean(dim=dim, keepdim=keepdim)
    return out


def knn(x, k):
    inner = -2 * torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    idx = pairwise_distance.topk(k=k, dim=-1)[1]  # (batch_size, num_points, k)
    return idx


class NDF(nn.Module):
    def __init__(
        self,
        latent_dim,
        sigmoid=True,
        return_features=False,
        acts="all",
        scaling=10.0,
    ):
        super().__init__()

        self.latent_dim = latent_dim
        self.scaling = (
            scaling  # scaling up the point cloud/query points to be larger helps
        )

        self.model_type = "pointnet"
        self.encoder = VNN_ResnetPointnet(c_dim=latent_dim)

    def forward(self, input):
        enc_in = input * self.scaling
        enc_out = self.encoder(enc_in)
        enc_out = torch.flatten(enc_out, start_dim=1)
        return enc_out


class VNN_ResnetPointnet(nn.Module):
    """
    Args:
        c_dim (int): dimension of latent code c
        dim (int): input points dimension
        hidden_dim (int): hidden dimension of the network
    """

    def __init__(self, c_dim=128, dim=3, hidden_dim=128, k=20, meta_output=None):
        super().__init__()
        self.c_dim = c_dim
        self.k = k
        self.meta_output = meta_output

        self.conv_pos = VNLinearLeakyReLU(
            3, 128, negative_slope=0.2, share_nonlinearity=False, use_batchnorm=False
        )
        self.fc_pos = VNLinear(128, 2 * hidden_dim)
        self.block_0 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_1 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_2 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_3 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_4 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.fc_c = VNLinear(hidden_dim, c_dim)

        self.actvn_c = VNLeakyReLU(
            hidden_dim, negative_slope=0.2, share_nonlinearity=False
        )
        self.pool = meanpool

        if meta_output == "invariant_latent":
            self.std_feature = VNStdFeature(
                c_dim, dim=3, normalize_frame=True, use_batchnorm=False
            )
        elif meta_output == "invariant_latent_linear":
            self.std_feature = VNStdFeature(
                c_dim, dim=3, normalize_frame=True, use_batchnorm=False
            )
            self.vn_inv = VNLinear(c_dim, 3)
        elif meta_output == "equivariant_latent_linear":
            self.vn_inv = VNLinear(c_dim, 3)

    def forward(self, p):
        batch_size = p.size(0)
        p = p.unsqueeze(1).transpose(2, 3)
        # mean = get_graph_mean(p, k=self.k)
        # mean = p_trans.mean(dim=-1, keepdim=True).expand(p_trans.size())
        feat = self.get_graph_feature_cross(p, k=self.k)
        net = self.conv_pos(feat)
        net = self.pool(net, dim=-1)

        net = self.fc_pos(net)

        net = self.block_0(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_1(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_2(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_3(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_4(net)

        # Recude to  B x F
        net = self.pool(net, dim=-1)

        c = self.fc_c(self.actvn_c(net))

        if self.meta_output == "invariant_latent":
            c_std, z0 = self.std_feature(c)
            return c, c_std
        elif self.meta_output == "invariant_latent_linear":
            c_std, z0 = self.std_feature(c)
            c_std = self.vn_inv(c_std)
            return c, c_std
        elif self.meta_output == "equivariant_latent_linear":
            c_std = self.vn_inv(c)
            return c, c_std

        return c

    def get_graph_feature_cross(self, x, k=20, idx=None):
        batch_size = x.size(0)
        num_points = x.size(3)
        x = x.view(batch_size, -1, num_points)
        if idx is None:
            idx = knn(x, k=k)  # (batch_size, num_points, k)
        # device = torch.device("cuda") if USE_CUDA else torch.device("cpu")
        device = x.device

        idx_base = (
            torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
        )

        idx = idx + idx_base

        idx = idx.view(-1)

        _, num_dims, _ = x.size()
        num_dims = num_dims // 3

        x = x.transpose(
            2, 1
        ).contiguous()  # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
        feature = x.view(batch_size * num_points, -1)[idx, :]
        feature = feature.view(batch_size, num_points, k, num_dims, 3)
        x = x.view(batch_size, num_points, 1, num_dims, 3).repeat(1, 1, k, 1, 1)
        cross = torch.cross(feature, x, dim=-1)

        feature = (
            torch.cat((feature - x, x, cross), dim=3)
            .permute(0, 3, 4, 1, 2)
            .contiguous()
        )

        return feature


class VNLinearLeakyReLU(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        dim=5,
        share_nonlinearity=False,
        use_batchnorm=True,
        negative_slope=0.2,
    ):
        super(VNLinearLeakyReLU, self).__init__()
        self.dim = dim
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm
        self.negative_slope = negative_slope

        # Conv
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

        # BatchNorm
        self.use_batchnorm = use_batchnorm
        if use_batchnorm == True:
            self.batchnorm = VNBatchNorm(out_channels, dim=dim)

        # LeakyReLU
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, out_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        # Conv
        p = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        # InstanceNorm
        if self.use_batchnorm == True:
            p = self.batchnorm(p)
        # LeakyReLU
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (p * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * p + (1 - self.negative_slope) * (
            mask * p + (1 - mask) * (p - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNBatchNorm(nn.Module):
    def __init__(self, num_features, dim):
        super(VNBatchNorm, self).__init__()
        self.dim = dim
        if dim == 3 or dim == 4:
            self.bn = nn.BatchNorm1d(num_features)
        elif dim == 5:
            self.bn = nn.BatchNorm2d(num_features)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        norm = torch.sqrt((x * x).sum(2))
        norm_bn = self.bn(norm)
        norm = norm.unsqueeze(2)
        norm_bn = norm_bn.unsqueeze(2)
        x = x / norm * norm_bn

        return x


class VNLinear(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(VNLinear, self).__init__()
        self.map_to_feat = nn.Linear(in_channels, out_channels, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        x_out = self.map_to_feat(x.transpose(1, -1)).transpose(1, -1)
        return x_out


class VNLeakyReLU(nn.Module):
    def __init__(self, in_channels, share_nonlinearity=False, negative_slope=0.2):
        super(VNLeakyReLU, self).__init__()
        if share_nonlinearity == True:
            self.map_to_dir = nn.Linear(in_channels, 1, bias=False)
        else:
            self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)
        self.negative_slope = negative_slope

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        d = self.map_to_dir(x.transpose(1, -1)).transpose(1, -1)
        dotprod = (x * d).sum(2, keepdim=True)
        mask = (dotprod >= 0).float()
        d_norm_sq = (d * d).sum(2, keepdim=True)
        x_out = self.negative_slope * x + (1 - self.negative_slope) * (
            mask * x + (1 - mask) * (x - (dotprod / (d_norm_sq + EPS)) * d)
        )
        return x_out


class VNStdFeature(nn.Module):
    def __init__(
        self,
        in_channels,
        dim=4,
        normalize_frame=False,
        share_nonlinearity=False,
        use_batchnorm=True,
    ):
        super(VNStdFeature, self).__init__()
        self.dim = dim
        self.normalize_frame = normalize_frame
        self.share_nonlinearity = share_nonlinearity
        self.use_batchnorm = use_batchnorm

        self.vn1 = VNLinearLeakyReLU(
            in_channels,
            in_channels // 2,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        self.vn2 = VNLinearLeakyReLU(
            in_channels // 2,
            in_channels // 4,
            dim=dim,
            share_nonlinearity=share_nonlinearity,
            use_batchnorm=use_batchnorm,
        )
        if normalize_frame:
            self.vn_lin = nn.Linear(in_channels // 4, 2, bias=False)
        else:
            self.vn_lin = nn.Linear(in_channels // 4, 3, bias=False)

    def forward(self, x):
        """
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        """
        z0 = self.vn1(x)
        z0 = self.vn2(z0)
        z0 = self.vn_lin(z0.transpose(1, -1)).transpose(1, -1)

        if self.normalize_frame:
            # make z0 orthogonal. u2 = v2 - proj_u1(v2)
            v1 = z0[:, 0, :]
            # u1 = F.normalize(v1, dim=1)
            v1_norm = torch.sqrt((v1 * v1).sum(1, keepdim=True))
            u1 = v1 / (v1_norm + EPS)
            v2 = z0[:, 1, :]
            v2 = v2 - (v2 * u1).sum(1, keepdim=True) * u1
            # u2 = F.normalize(u2, dim=1)
            v2_norm = torch.sqrt((v2 * v2).sum(1, keepdim=True))
            u2 = v2 / (v2_norm + EPS)

            # compute the cross product of the two output vectors
            u3 = torch.cross(u1, u2)
            z0 = torch.stack([u1, u2, u3], dim=1).transpose(1, 2)
        else:
            z0 = z0.transpose(1, 2)

        if self.dim == 4:
            x_std = torch.einsum("bijm,bjkm->bikm", x, z0)
        elif self.dim == 3:
            x_std = torch.einsum("bij,bjk->bik", x, z0)
        elif self.dim == 5:
            x_std = torch.einsum("bijmn,bjkmn->bikmn", x, z0)

        return x_std, z0


# Resnet Blocks
class VNResnetBlockFC(nn.Module):
    """Fully connected ResNet Block class.

    Args:
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
    """

    def __init__(self, size_in, size_out=None, size_h=None):
        super().__init__()
        # Attributes
        if size_out is None:
            size_out = size_in

        if size_h is None:
            size_h = min(size_in, size_out)

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        self.fc_0 = VNLinear(size_in, size_h)
        self.fc_1 = VNLinear(size_h, size_out)
        self.actvn_0 = VNLeakyReLU(
            size_in, negative_slope=0.2, share_nonlinearity=False
        )
        self.actvn_1 = VNLeakyReLU(size_h, negative_slope=0.2, share_nonlinearity=False)

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = VNLinear(size_in, size_out)
        # Initialization
        nn.init.zeros_(self.fc_1.map_to_feat.weight)

    def forward(self, x):
        net = self.fc_0(self.actvn_0(x))
        dx = self.fc_1(self.actvn_1(net))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx

NCF_policies/algo2/ncf/ndf/ndf.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation
# https://arxiv.org/abs/2112.05124
# Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B. Tenenbaum, Alberto Rodriguez, 
# Pulkit Agrawal, Vincent Sitzmann
# --------------------------------------------------------

import os, os.path as osp
import torch
import numpy as np
import trimesh
import random
import copy
import plotly.graph_objects as go

from ndf_robot.utils import torch_util, trimesh_util
from ndf_robot.utils.plotly_save import plot3d


class NDFAlignmentCheck:
    def __init__(self, model, pcd1, pcd2, model_type='pointnet', opt_iterations=500, sigma=0.025, trimesh_viz=False):
        self.model = model
        self.model_type = model_type
        self.opt_iterations = opt_iterations
        self.sigma = sigma
        self.trimesh_viz = trimesh_viz

        self.perturb_scale = 0.001
        self.perturb_decay = 0.5
        self.n_pts =  1500
        self.n_opt_pts = 500
        self.prepare_inputs(pcd1, pcd2)

        self.loss_fn = torch.nn.L1Loss()
        if torch.cuda.is_available():
            self.dev = torch.device('cuda:0')
        else:
            self.dev = torch.device('cpu')
        self.model = self.model.to(self.dev)
        self.model.eval()

        self.viz_path = 'visualization'
        if not osp.exists(self.viz_path):
            os.makedirs(self.viz_path)

        self.video_viz_path = 'vid_visualization'
        if not osp.exists(self.video_viz_path):
            os.makedirs(self.video_viz_path)

        self._cam_frame_scene_dict()

    def _cam_frame_scene_dict(self):
        self.cam_frame_scene_dict = {}
        cam_up_vec = [0, 1, 0]
        plotly_camera = {
            'up': {'x': cam_up_vec[0], 'y': cam_up_vec[1],'z': cam_up_vec[2]},
            'center': {'x': 0, 'y': 0, 'z': 0},
            'eye': {'x': -0.6, 'y': -0.6, 'z': 0.4},
        }

        plotly_scene = {
            'xaxis': 
                {
                    'backgroundcolor': 'rgb(255, 255, 255)',
                    'gridcolor': 'white',
                    'zerolinecolor': 'white',
                    'tickcolor': 'rgb(255, 255, 255)',
                    'showticklabels': False,
                    'showbackground': False,
                    'showaxeslabels': False,
                    'visible': False,
                    'range': [-0.5, 0.5]},
            'yaxis': 
                {
                    'backgroundcolor': 'rgb(255, 255, 255)',
                    'gridcolor': 'white',
                    'zerolinecolor': 'white',
                    'tickcolor': 'rgb(255, 255, 255)',
                    'showticklabels': False,
                    'showbackground': False,
                    'showaxeslabels': False,
                    'visible': False,
                    'range': [-0.5, 0.5]},
            'zaxis': 
                {
                    'backgroundcolor': 'rgb(255, 255, 255)',
                    'gridcolor': 'white',
                    'zerolinecolor': 'white',
                    'tickcolor': 'rgb(255, 255, 255)',
                    'showticklabels': False,
                    'showbackground': False,
                    'showaxeslabels': False,
                    'visible': False,
                    'range': [-0.5, 0.5]},
        }
        self.cam_frame_scene_dict['camera'] = plotly_camera
        self.cam_frame_scene_dict['scene'] = plotly_scene

    def plotly_create_local_frame(self, transform=None, length=0.03):
        if transform is None:
            transform = np.eye(4)

        x_vec = transform[:-1, 0] * length
        y_vec = transform[:-1, 1] * length
        z_vec = transform[:-1, 2] * length

        origin = transform[:-1, -1]

        lw = 8
        x_data = go.Scatter3d(
            x=[origin[0], x_vec[0] + origin[0]], y=[origin[1], x_vec[1] + origin[1]], z=[origin[2], x_vec[2] + origin[2]],
            line=dict(
                color='red',
                width=lw
            ),
            marker=dict(
                size=0.0
            )
        )
        y_data = go.Scatter3d(
            x=[origin[0], y_vec[0] + origin[0]], y=[origin[1], y_vec[1] + origin[1]], z=[origin[2], y_vec[2] + origin[2]],
            line=dict(
                color='green',
                width=lw
            ),
            marker=dict(
                size=0.0
            )
        )
        z_data = go.Scatter3d(
            x=[origin[0], z_vec[0] + origin[0]], y=[origin[1], z_vec[1] + origin[1]], z=[origin[2], z_vec[2] + origin[2]],
            line=dict(
                color='blue',
                width=lw
            ),
            marker=dict(
                size=0.0
            )
        )
        # fig = go.Figure(data=[x_data, y_data, z_data])
        # fig.show()

        data = [x_data, y_data, z_data]
        return data

    def prepare_inputs(self, pcd1, pcd2):
        pcd1 = pcd1 - np.mean(pcd1, axis=0)
        pcd2 = pcd2 - np.mean(pcd2, axis=0)
        self.pcd1 = pcd1
        self.pcd2 = pcd2

        tpcd1 = trimesh.PointCloud(self.pcd1[:self.n_pts])
        tpcd2 = trimesh.PointCloud(self.pcd2[:self.n_pts])
        # tpcd1.show()
        # tpcd2.show()

    def sample_pts(self, show_recon=False, return_scene=False, visualize_all_inits=False, render_video=False):
        # sample query points
        query_pts = np.random.normal(0.0, self.sigma, size=(self.n_opt_pts, 3))
        
        # put the query points at one of the points in the point cloud
        q_offset_ind = np.random.randint(self.pcd1.shape[0])
        q_offset = self.pcd1[q_offset_ind]
        q_offset *= 1.2
        reference_query_pts = query_pts + q_offset

        reference_model_input = {}
        ref_query_pts = torch.from_numpy(reference_query_pts[:self.n_opt_pts]).float().to(self.dev)
        ref_shape_pcd = torch.from_numpy(self.pcd1[:self.n_pts]).float().to(self.dev)
        reference_model_input['coords'] = ref_query_pts[None, :, :]
        reference_model_input['point_cloud'] = ref_shape_pcd[None, :, :]

        # get the descriptors for these reference query points
        reference_latent = self.model.extract_latent(reference_model_input).detach()
        reference_act_hat = self.model.forward_latent(reference_latent, reference_model_input['coords']).detach()

        # set up the optimization
        if 'dgcnn' in self.model_type:
            full_opt = 5   # dgcnn can't fit 10 initialization in memory
        else:
            full_opt = 10
        best_loss = np.inf
        best_tf = np.eye(4)
        best_idx = 0
        tf_list = []
        M = full_opt

        trans = (torch.rand((M, 3)) * 0.1).float().to(self.dev)
        rot = torch.rand(M, 3).float().to(self.dev)
        trans.requires_grad_()
        rot.requires_grad_()
        opt = torch.optim.Adam([trans, rot], lr=1e-2)

        rand_rot_init = (torch.rand((M, 3)) * 2*np.pi).float().to(self.dev)
        rand_mat_init = torch_util.angle_axis_to_rotation_matrix(rand_rot_init)
        rand_mat_init = rand_mat_init.squeeze().float().to(self.dev)

        # now randomly initialize a copy of the query points
        opt_query_pts = torch.from_numpy(query_pts).float().to(self.dev)
        opt_query_pts = opt_query_pts[:self.n_opt_pts][None, :, :].repeat((M, 1, 1))
        X = torch_util.transform_pcd_torch(opt_query_pts, rand_mat_init)

        opt_model_input = {}
        opt_model_input['coords'] = X

        mi_point_cloud = []
        for ii in range(M):
            mi_point_cloud.append(torch.from_numpy(self.pcd2[:self.n_pts]).float().to(self.dev))
        mi_point_cloud = torch.stack(mi_point_cloud, 0)
        opt_model_input['point_cloud'] = mi_point_cloud
        opt_latent = self.model.extract_latent(opt_model_input).detach()

        loss_values = []
        vid_plot_idx = None  # we will set this during the optimization

        # run optimization
        pcd_traj_list = {}
        for jj in range(M):
            pcd_traj_list[jj] = []
            pcd_traj_list[jj].append(np.mean(X[jj].detach().cpu().numpy(), axis=0))
        for i in range(self.opt_iterations):
            T_mat = torch_util.angle_axis_to_rotation_matrix(rot).squeeze()
            noise_vec = (torch.randn(X.size()) * (self.perturb_scale / ((i+1)**(self.perturb_decay)))).to(self.dev)
            X_perturbed = X + noise_vec
            X_new = torch_util.transform_pcd_torch(X_perturbed, T_mat) + trans[:, None, :].repeat((1, X.size(1), 1))

            ######################### stuff for visualizing the reconstruction ##################33

            for jj in range(M):
                X_np = X_new[jj].detach().cpu().numpy()
                centroid = np.mean(X_np, axis=0)
                pcd_traj_list[jj].append(centroid)

            if i == 0 and show_recon:
                jj = 0
                shape_mi = {}
                shape_mi['point_cloud'] = opt_model_input['point_cloud'][jj][None, :, :].detach()
                shape_np = shape_mi['point_cloud'].cpu().numpy().squeeze()
                shape_mean = np.mean(shape_np, axis=0)
                inliers = np.where(np.linalg.norm(shape_np - shape_mean, 2, 1) < 0.2)[0]
                shape_np = shape_np[inliers]
                shape_pcd = trimesh.PointCloud(shape_np)
                bb = shape_pcd.bounding_box
                bb_scene = trimesh.Scene(); bb_scene.add_geometry([shape_pcd, bb]) 
                # bb_scene.show()

                eval_pts = bb.sample_volume(10000)
                shape_mi['coords'] = torch.from_numpy(eval_pts)[None, :, :].float().to(self.dev).detach()
                out = self.model(shape_mi)
                thresh = 0.1
                in_inds = torch.where(out['occ'].squeeze() > thresh)[0].cpu().numpy()
                out_inds = torch.where(out['occ'].squeeze() < thresh)[0].cpu().numpy()

                in_pts = eval_pts[in_inds]
                out_pts = eval_pts[out_inds]
                if self.trimesh_viz:
                    scene = trimesh_util.trimesh_show([in_pts])
                    in_scene = trimesh_util.trimesh_show([in_pts, shape_np])
                self._cam_frame_scene_dict()
                plot3d(
                    [in_pts, shape_np],
                    ['blue', 'black'], 
                    osp.join(self.viz_path, 'recon_overlay_test.html'),
                    scene_dict=self.cam_frame_scene_dict,
                    z_plane=False)

            ###############################################################################

            act_hat = self.model.forward_latent(opt_latent, X_new)
            t_size = reference_act_hat.size()

            losses = [self.loss_fn(act_hat[ii].view(t_size), reference_act_hat) for ii in range(M)]

            loss = torch.mean(torch.stack(losses))
            if i % 100 == 0:
                losses_str = ['%f' % val.item() for val in losses]
                loss_str = ', '.join(losses_str)
                print(f'i: {i}, losses: {loss_str}')
            loss_values.append(loss.item())
            opt.zero_grad()
            loss.backward()
            opt.step()

            if i > 5 and (vid_plot_idx is None):
                # try to guess which run in the batch will lead to lowest cost
                vid_plot_idx = torch.argmin(torch.stack(losses)).item()
                print('vid plot idx: ', vid_plot_idx)
                plot_rand_mat_np = rand_mat_init[vid_plot_idx].detach().cpu().numpy() 
            
            if i < 200:
                render_iter = (i % 4 == 0)
            else:
                render_iter = (i % 8 == 0)

            if render_video and render_iter and (vid_plot_idx is not None):
                # save image for each iteration
                transform_mat_np = T_mat[vid_plot_idx].detach().cpu().numpy()
                transform_mat_np[:-1, -1] = trans[vid_plot_idx].detach().cpu().numpy()
                frame_tf = np.matmul(transform_mat_np, plot_rand_mat_np)

                frame_data = self.plotly_create_local_frame(transform=frame_tf)
                plot3d(
                    [self.pcd2, X_new[vid_plot_idx].detach().cpu().numpy(), np.asarray(pcd_traj_list[vid_plot_idx])],
                    ['blue', 'black', 'red'], 
                    osp.join(self.video_viz_path, 'opt_iter_%d.png' % i),
                    scene_dict=self.cam_frame_scene_dict,
                    z_plane=False,
                    extra_data=frame_data)

        best_idx = torch.argmin(torch.stack(losses)).item()
        best_loss = losses[best_idx]
        print('best loss: %f, best_idx: %d' % (best_loss, best_idx))

        best_X = X_new[best_idx].detach().cpu().numpy()

        offset = np.array([0.4, 0, 0])
        vpcd1 = copy.deepcopy(self.pcd1)
        vquery1 = copy.deepcopy(reference_query_pts)

        vpcd1 += offset
        vquery1 += offset

        trans_best, rot_best = trans[best_idx], rot[best_idx]
        transform_mat_np = torch_util.angle_axis_to_rotation_matrix(rot_best.view(1, -1)).squeeze().detach().cpu().numpy()
        transform_mat_np[:-1, -1] = trans_best.detach().cpu().numpy()
        rand_mat_np = rand_mat_init[best_idx].detach().cpu().numpy() 

        frame1_tf = np.eye(4)
        frame1_tf[:-1, -1] = (q_offset + offset)
        frame2_tf = np.matmul(transform_mat_np, rand_mat_np)

        frame1 = self.plotly_create_local_frame(transform=frame1_tf)
        frame2 = self.plotly_create_local_frame(transform=rand_mat_np)
        frame3 = self.plotly_create_local_frame(transform=frame2_tf)
        frame_data = frame1 + frame2 + frame3
        plot3d(
            [vpcd1, vquery1, self.pcd2, best_X, np.asarray(pcd_traj_list[best_idx])],
            ['purple', 'black', 'blue', 'black', 'red'], 
            osp.join(self.viz_path, 'best_alignment.html'),
            scene_dict=self.cam_frame_scene_dict,
            z_plane=False,
            extra_data=frame_data)

        best_scene = None
        if self.trimesh_viz:
            local_frame_1 = trimesh.creation.axis(origin_size=0.002, transform=frame1_tf, origin_color=None, axis_radius=None, axis_length=0.03)
            local_frame_2 = trimesh.creation.axis(origin_size=0.002, transform=rand_mat_np, origin_color=None, axis_radius=None, axis_length=0.03)
            local_frame_3 = trimesh.creation.axis(origin_size=0.002, transform=frame2_tf, origin_color=None, axis_radius=None, axis_length=0.03)
             
            best_scene = trimesh_util.trimesh_show([vpcd1, vquery1 , self.pcd2, best_X, pcd_traj_list[best_idx]], show=False)
            best_scene.add_geometry([local_frame_1, local_frame_2, local_frame_3])
            best_scene.show()

            # all the runs from different initializations
        if visualize_all_inits:
            for j in range(M):
                trans_j, rot_j = trans[j], rot[j]
                transform_mat_np = torch_util.angle_axis_to_rotation_matrix(rot_j.view(1, -1)).squeeze().detach().cpu().numpy()
                transform_mat_np[:-1, -1] = trans_j.detach().cpu().numpy()

                rand_mat_np = rand_mat_init[j].detach().cpu().numpy() 

                frame1_tf = np.eye(4)
                frame1_tf[:-1, -1] = (q_offset + offset)
                frame2_tf = np.matmul(transform_mat_np, rand_mat_np)

                X_show = X_new[j].detach().cpu().numpy()

                frame1 = self.plotly_create_local_frame(transform=frame1_tf)
                frame2 = self.plotly_create_local_frame(transform=rand_mat_np)
                frame3 = self.plotly_create_local_frame(transform=frame2_tf)
                frame_data = frame1 + frame2 + frame3
                plot3d(
                    [vpcd1, vquery1, self.pcd2, X_show, np.asarray(pcd_traj_list[j])],
                    ['purple', 'black', 'blue', 'black', 'red'], 
                    osp.join(self.viz_path, 'alignment_%d.html' % j),
                    scene_dict=self.cam_frame_scene_dict,
                    z_plane=False,
                    extra_data=frame_data)

                if self.trimesh_viz:
                    local_frame_1 = trimesh.creation.axis(origin_size=0.002, transform=frame1_tf, origin_color=None, axis_radius=None, axis_length=0.03)
                    local_frame_2 = trimesh.creation.axis(origin_size=0.002, transform=rand_mat_np, origin_color=None, axis_radius=None, axis_length=0.03)
                    local_frame_3 = trimesh.creation.axis(origin_size=0.002, transform=frame2_tf, origin_color=None, axis_radius=None, axis_length=0.03)
                    
                    scene = trimesh_util.trimesh_show([vpcd1, vquery1 , self.pcd2, X_show, pcd_traj_list[j]], show=False)
                    scene.add_geometry([local_frame_1, local_frame_2, local_frame_3])
                    scene.show()
        if return_scene:
            return best_scene
NCF_policies/algo2/ncf/ndf/ndf_alignment.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation
# https://arxiv.org/abs/2112.05124
# Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B. Tenenbaum, Alberto Rodriguez, 
# Pulkit Agrawal, Vincent Sitzmann
# --------------------------------------------------------

import torch
import torch.nn as nn
from algo.ncf.ndf.layers_equi import *


def maxpool(x, dim=-1, keepdim=False):
    out, _ = x.max(dim=dim, keepdim=keepdim)
    return out


def meanpool(x, dim=-1, keepdim=False):
    out = x.mean(dim=dim, keepdim=keepdim)
    return out


class VNN_DGCNN(nn.Module):
    def __init__(self, c_dim=128, dim=3, hidden_dim=64, k=20):
        super(VNN_DGCNN, self).__init__()
        self.c_dim = c_dim
        self.k = k

        self.conv1 = VNLinearLeakyReLU(2, hidden_dim)
        self.conv2 = VNLinearLeakyReLU(hidden_dim * 2, hidden_dim)
        self.conv3 = VNLinearLeakyReLU(hidden_dim * 2, hidden_dim)
        self.conv4 = VNLinearLeakyReLU(hidden_dim * 2, hidden_dim)

        self.pool1 = meanpool
        self.pool2 = meanpool
        self.pool3 = meanpool
        self.pool4 = meanpool

        self.conv_c = VNLinearLeakyReLU(
            hidden_dim * 4, c_dim, dim=4, share_nonlinearity=True
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = x.unsqueeze(1).transpose(2, 3)
        x = get_graph_feature(x, k=self.k)
        x = self.conv1(x)
        x1 = self.pool1(x)

        x = get_graph_feature(x1, k=self.k)
        x = self.conv2(x)
        x2 = self.pool2(x)

        x = get_graph_feature(x2, k=self.k)
        x = self.conv3(x)
        x3 = self.pool3(x)

        x = get_graph_feature(x3, k=self.k)
        x = self.conv4(x)
        x4 = self.pool4(x)

        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.conv_c(x)
        x = x.mean(dim=-1, keepdim=False)

        return x


class VNNOccNet(nn.Module):
    def __init__(
        self,
        latent_dim,
        model_type="pointnet",
        sigmoid=True,
        return_features=False,
        acts="all",
        scaling=10.0,
    ):
        super().__init__()

        self.latent_dim = latent_dim
        self.scaling = (
            scaling  # scaling up the point cloud/query points to be larger helps
        )
        self.return_features = return_features

        if model_type == "dgcnn":
            self.model_type = "dgcnn"
            self.encoder = VNN_DGCNN(c_dim=latent_dim)  # modified resnet-18
        else:
            self.model_type = "pointnet"
            self.encoder = VNN_ResnetPointnet(c_dim=latent_dim)  # modified resnet-18

        self.decoder = DecoderInner(
            dim=3,
            z_dim=latent_dim,
            c_dim=0,
            hidden_size=latent_dim,
            leaky=True,
            sigmoid=sigmoid,
            return_features=return_features,
            acts=acts,
        )

    def forward(self, input):
        out_dict = {}

        enc_in = input["point_cloud"] * self.scaling
        query_points = input["coords"] * self.scaling

        z = self.encoder(enc_in)

        if self.return_features:
            out_dict["occ"], out_dict["features"] = self.decoder(query_points, z)
        else:
            out_dict["occ"] = self.decoder(query_points, z)

        return out_dict

    def extract_latent(self, input):
        enc_in = input["point_cloud"] * self.scaling
        z = self.encoder(enc_in)
        return z

    def forward_latent(self, z, coords):
        out_dict = {}
        coords = coords * self.scaling
        out_dict["occ"], out_dict["features"] = self.decoder(coords, z)

        return out_dict["features"]


class VNN_ResnetPointnet(nn.Module):
    """DGCNN-based VNN encoder network with ResNet blocks.

    Args:
        c_dim (int): dimension of latent code c
        dim (int): input points dimension
        hidden_dim (int): hidden dimension of the network
    """

    def __init__(self, c_dim=128, dim=3, hidden_dim=128, k=20, meta_output=None):
        super().__init__()
        self.c_dim = c_dim
        self.k = k
        self.meta_output = meta_output

        self.conv_pos = VNLinearLeakyReLU(
            3, 128, negative_slope=0.2, share_nonlinearity=False, use_batchnorm=False
        )
        self.fc_pos = VNLinear(128, 2 * hidden_dim)
        self.block_0 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_1 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_2 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_3 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.block_4 = VNResnetBlockFC(2 * hidden_dim, hidden_dim)
        self.fc_c = VNLinear(hidden_dim, c_dim)

        self.actvn_c = VNLeakyReLU(
            hidden_dim, negative_slope=0.2, share_nonlinearity=False
        )
        self.pool = meanpool

        if meta_output == "invariant_latent":
            self.std_feature = VNStdFeature(
                c_dim, dim=3, normalize_frame=True, use_batchnorm=False
            )
        elif meta_output == "invariant_latent_linear":
            self.std_feature = VNStdFeature(
                c_dim, dim=3, normalize_frame=True, use_batchnorm=False
            )
            self.vn_inv = VNLinear(c_dim, 3)
        elif meta_output == "equivariant_latent_linear":
            self.vn_inv = VNLinear(c_dim, 3)

    def forward(self, p):
        batch_size = p.size(0)
        p = p.unsqueeze(1).transpose(2, 3)
        # mean = get_graph_mean(p, k=self.k)
        # mean = p_trans.mean(dim=-1, keepdim=True).expand(p_trans.size())
        feat = get_graph_feature_cross(p, k=self.k)
        net = self.conv_pos(feat)
        net = self.pool(net, dim=-1)

        net = self.fc_pos(net)

        net = self.block_0(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_1(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_2(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_3(net)
        pooled = self.pool(net, dim=-1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=1)

        net = self.block_4(net)

        # Recude to  B x F
        net = self.pool(net, dim=-1)

        c = self.fc_c(self.actvn_c(net))

        if self.meta_output == "invariant_latent":
            c_std, z0 = self.std_feature(c)
            return c, c_std
        elif self.meta_output == "invariant_latent_linear":
            c_std, z0 = self.std_feature(c)
            c_std = self.vn_inv(c_std)
            return c, c_std
        elif self.meta_output == "equivariant_latent_linear":
            c_std = self.vn_inv(c)
            return c, c_std

        return c


class DecoderInner(nn.Module):
    """Decoder class.

    It does not perform any form of normalization.

    Args:
        dim (int): input dimension
        z_dim (int): dimension of latent code z
        c_dim (int): dimension of latent conditioned code c
        hidden_size (int): hidden size of Decoder network
        leaky (bool): whether to use leaky ReLUs
    """

    def __init__(
        self,
        dim=3,
        z_dim=128,
        c_dim=128,
        hidden_size=128,
        leaky=False,
        return_features=False,
        sigmoid=True,
        acts="all",
    ):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim

        self.acts = acts
        if self.acts not in ["all", "inp", "first_rn", "inp_first_rn"]:
            # self.acts = 'all'
            raise ValueError(
                'Please provide "acts" equal to one of the following: "all", "inp", "first_rn", "inp_first_rn"'
            )

        # Submodules
        if z_dim > 0:
            self.z_in = VNLinear(z_dim, z_dim)
        if c_dim > 0:
            self.c_in = VNLinear(c_dim, c_dim)

        self.fc_in = nn.Linear(z_dim * 2 + c_dim * 2 + 1, hidden_size)

        self.block0 = ResnetBlockFC(hidden_size)
        self.block1 = ResnetBlockFC(hidden_size)
        self.block2 = ResnetBlockFC(hidden_size)
        self.block3 = ResnetBlockFC(hidden_size)
        self.block4 = ResnetBlockFC(hidden_size)
        self.return_features = return_features

        self.fc_out = nn.Linear(hidden_size, 1)
        self.sigmoid = sigmoid

        if not leaky:
            self.actvn = F.relu
        else:
            self.actvn = lambda x: F.leaky_relu(x, 0.2)

    def forward(self, p, z, c=None, **kwargs):
        batch_size, T, D = p.size()
        acts = []
        acts_inp = []
        acts_first_rn = []
        acts_inp_first_rn = []

        if isinstance(c, tuple):
            c, c_meta = c

        net = (p * p).sum(2, keepdim=True)

        if self.z_dim != 0:
            z = z.view(batch_size, -1, D).contiguous()
            net_z = torch.einsum("bmi,bni->bmn", p, z)
            z_dir = self.z_in(z)
            z_inv = (z * z_dir).sum(-1).unsqueeze(1).repeat(1, T, 1)
            net = torch.cat([net, net_z, z_inv], dim=2)

        if self.c_dim != 0:
            c = c.view(batch_size, -1, D).contiguous()
            net_c = torch.einsum("bmi,bni->bmn", p, c)
            c_dir = self.c_in(c)
            c_inv = (c * c_dir).sum(-1).unsqueeze(1).repeat(1, T, 1)
            net = torch.cat([net, net_c, c_inv], dim=2)

        acts.append(net)
        acts_inp.append(net)
        acts_inp_first_rn.append(net)

        net = self.fc_in(net)
        acts.append(net)
        # acts_inp.append(net)
        # acts_inp_first_rn.append(net)

        net = self.block0(net)
        acts.append(net)
        # acts_inp_first_rn.append(net)
        acts_first_rn.append(net)

        net = self.block1(net)
        acts.append(net)
        net = self.block2(net)
        acts.append(net)
        net = self.block3(net)
        acts.append(net)
        net = self.block4(net)
        last_act = net
        acts.append(net)

        out = self.fc_out(self.actvn(net))
        out = out.squeeze(-1)

        if self.sigmoid:
            out = F.sigmoid(out)

        if self.return_features:
            # acts = torch.cat(acts, dim=-1)
            if self.acts == "all":
                acts = torch.cat(acts, dim=-1)
            elif self.acts == "inp":
                acts = torch.cat(acts_inp, dim=-1)
            elif self.acts == "last":
                acts = last_act
            elif self.acts == "inp_first_rn":
                acts = torch.cat(acts_inp_first_rn, dim=-1)
            acts = F.normalize(acts, p=2, dim=-1)
            return out, acts
        else:
            return out


class DecoderCBatchNorm(nn.Module):
    """Decoder with conditional batch normalization (CBN) class.

    Args:
        dim (int): input dimension
        z_dim (int): dimension of latent code z
        c_dim (int): dimension of latent conditioned code c
        hidden_size (int): hidden size of Decoder network
        leaky (bool): whether to use leaky ReLUs
        legacy (bool): whether to use the legacy structure
    """

    def __init__(
        self, dim=3, z_dim=128, c_dim=128, hidden_size=256, leaky=False, legacy=False
    ):
        super().__init__()
        self.z_dim = z_dim
        if not z_dim == 0:
            self.fc_z = nn.Linear(z_dim, hidden_size)

        self.fc_p = nn.Conv1d(dim, hidden_size, 1)
        self.block0 = CResnetBlockConv1d(c_dim, hidden_size, legacy=legacy)
        self.block1 = CResnetBlockConv1d(c_dim, hidden_size, legacy=legacy)
        self.block2 = CResnetBlockConv1d(c_dim, hidden_size, legacy=legacy)
        self.block3 = CResnetBlockConv1d(c_dim, hidden_size, legacy=legacy)
        self.block4 = CResnetBlockConv1d(c_dim, hidden_size, legacy=legacy)

        if not legacy:
            self.bn = CBatchNorm1d(c_dim, hidden_size)
        else:
            self.bn = CBatchNorm1d_legacy(c_dim, hidden_size)

        self.fc_out = nn.Conv1d(hidden_size, 1, 1)

        if not leaky:
            self.actvn = F.relu
        else:
            self.actvn = lambda x: F.leaky_relu(x, 0.2)

    def forward(self, p, z, c, **kwargs):
        p = p.transpose(1, 2)
        batch_size, D, T = p.size()
        net = self.fc_p(p)

        if self.z_dim != 0:
            net_z = self.fc_z(z).unsqueeze(2)
            net = net + net_z

        net = self.block0(net, c)
        net = self.block1(net, c)
        net = self.block2(net, c)
        net = self.block3(net, c)
        net = self.block4(net, c)

        out = self.fc_out(self.actvn(self.bn(net, c)))
        out = out.squeeze(1)

        return out


class DecoderCBatchNorm2(nn.Module):
    """Decoder with CBN class 2.

    It differs from the previous one in that the number of blocks can be
    chosen.

    Args:
        dim (int): input dimension
        z_dim (int): dimension of latent code z
        c_dim (int): dimension of latent conditioned code c
        hidden_size (int): hidden size of Decoder network
        leaky (bool): whether to use leaky ReLUs
        n_blocks (int): number of ResNet blocks
    """

    def __init__(self, dim=3, z_dim=0, c_dim=128, hidden_size=256, n_blocks=5):
        super().__init__()
        self.z_dim = z_dim
        if z_dim != 0:
            self.fc_z = nn.Linear(z_dim, c_dim)

        self.conv_p = nn.Conv1d(dim, hidden_size, 1)
        self.blocks = nn.ModuleList(
            [CResnetBlockConv1d(c_dim, hidden_size) for i in range(n_blocks)]
        )

        self.bn = CBatchNorm1d(c_dim, hidden_size)
        self.conv_out = nn.Conv1d(hidden_size, 1, 1)
        self.actvn = nn.ReLU()

    def forward(self, p, z, c, **kwargs):
        p = p.transpose(1, 2)
        batch_size, D, T = p.size()
        net = self.conv_p(p)

        if self.z_dim != 0:
            c = c + self.fc_z(z)

        for block in self.blocks:
            net = block(net, c)

        out = self.conv_out(self.actvn(self.bn(net, c)))
        out = out.squeeze(1)

        return out


class DecoderCBatchNormNoResnet(nn.Module):
    """Decoder CBN with no ResNet blocks class.

    Args:
        dim (int): input dimension
        z_dim (int): dimension of latent code z
        c_dim (int): dimension of latent conditioned code c
        hidden_size (int): hidden size of Decoder network
        leaky (bool): whether to use leaky ReLUs
    """

    def __init__(self, dim=3, z_dim=128, c_dim=128, hidden_size=256, leaky=False):
        super().__init__()
        self.z_dim = z_dim
        if not z_dim == 0:
            self.fc_z = nn.Linear(z_dim, hidden_size)

        self.fc_p = nn.Conv1d(dim, hidden_size, 1)
        self.fc_0 = nn.Conv1d(hidden_size, hidden_size, 1)
        self.fc_1 = nn.Conv1d(hidden_size, hidden_size, 1)
        self.fc_2 = nn.Conv1d(hidden_size, hidden_size, 1)
        self.fc_3 = nn.Conv1d(hidden_size, hidden_size, 1)
        self.fc_4 = nn.Conv1d(hidden_size, hidden_size, 1)

        self.bn_0 = CBatchNorm1d(c_dim, hidden_size)
        self.bn_1 = CBatchNorm1d(c_dim, hidden_size)
        self.bn_2 = CBatchNorm1d(c_dim, hidden_size)
        self.bn_3 = CBatchNorm1d(c_dim, hidden_size)
        self.bn_4 = CBatchNorm1d(c_dim, hidden_size)
        self.bn_5 = CBatchNorm1d(c_dim, hidden_size)

        self.fc_out = nn.Conv1d(hidden_size, 1, 1)

        if not leaky:
            self.actvn = F.relu
        else:
            self.actvn = lambda x: F.leaky_relu(x, 0.2)

    def forward(self, p, z, c, **kwargs):
        p = p.transpose(1, 2)
        batch_size, D, T = p.size()
        net = self.fc_p(p)

        if self.z_dim != 0:
            net_z = self.fc_z(z).unsqueeze(2)
            net = net + net_z

        net = self.actvn(self.bn_0(net, c))
        net = self.fc_0(net)
        net = self.actvn(self.bn_1(net, c))
        net = self.fc_1(net)
        net = self.actvn(self.bn_2(net, c))
        net = self.fc_2(net)
        net = self.actvn(self.bn_3(net, c))
        net = self.fc_3(net)
        net = self.actvn(self.bn_4(net, c))
        net = self.fc_4(net)
        net = self.actvn(self.bn_5(net, c))
        out = self.fc_out(net)
        out = out.squeeze(1)

        return out


class DecoderBatchNorm(nn.Module):
    """Decoder with batch normalization class.

    Args:
        dim (int): input dimension
        z_dim (int): dimension of latent code z
        c_dim (int): dimension of latent conditioned code c
        hidden_size (int): hidden size of Decoder network
        leaky (bool): whether to use leaky ReLUs
    """

    def __init__(self, dim=3, z_dim=128, c_dim=128, hidden_size=256, leaky=False):
        super().__init__()
        self.z_dim = z_dim
        self.c_dim = c_dim

        # Submodules
        if not z_dim == 0:
            self.fc_z = nn.Linear(z_dim, hidden_size)

        if self.c_dim != 0:
            self.fc_c = nn.Linear(c_dim, hidden_size)
        self.fc_p = nn.Conv1d(dim, hidden_size, 1)
        self.block0 = ResnetBlockConv1d(hidden_size)
        self.block1 = ResnetBlockConv1d(hidden_size)
        self.block2 = ResnetBlockConv1d(hidden_size)
        self.block3 = ResnetBlockConv1d(hidden_size)
        self.block4 = ResnetBlockConv1d(hidden_size)

        self.bn = nn.BatchNorm1d(hidden_size)

        self.fc_out = nn.Conv1d(hidden_size, 1, 1)

        if not leaky:
            self.actvn = F.relu
        else:
            self.actvn = lambda x: F.leaky_relu(x, 0.2)

    def forward(self, p, z, c, **kwargs):
        p = p.transpose(1, 2)
        batch_size, D, T = p.size()
        net = self.fc_p(p)

        if self.z_dim != 0:
            net_z = self.fc_z(z).unsqueeze(2)
            net = net + net_z

        if self.c_dim != 0:
            net_c = self.fc_c(c).unsqueeze(2)
            net = net + net_c

        net = self.block0(net)
        net = self.block1(net)
        net = self.block2(net)
        net = self.block3(net)
        net = self.block4(net)

        out = self.fc_out(self.actvn(self.bn(net)))
        out = out.squeeze(1)

        return out


# Resnet Blocks
class ResnetBlockFC(nn.Module):
    """Fully connected ResNet Block class.

    Args:
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
    """

    def __init__(self, size_in, size_out=None, size_h=None):
        super().__init__()
        # Attributes
        if size_out is None:
            size_out = size_in

        if size_h is None:
            size_h = min(size_in, size_out)

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        self.fc_0 = nn.Linear(size_in, size_h)
        self.fc_1 = nn.Linear(size_h, size_out)
        self.actvn = nn.ReLU()

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = nn.Linear(size_in, size_out, bias=False)
        # Initialization
        nn.init.zeros_(self.fc_1.weight)

    def forward(self, x):
        net = self.fc_0(self.actvn(x))
        dx = self.fc_1(self.actvn(net))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx


class CResnetBlockConv1d(nn.Module):
    """Conditional batch normalization-based Resnet block class.

    Args:
        c_dim (int): dimension of latend conditioned code c
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
        norm_method (str): normalization method
        legacy (bool): whether to use legacy blocks
    """

    def __init__(
        self,
        c_dim,
        size_in,
        size_h=None,
        size_out=None,
        norm_method="batch_norm",
        legacy=False,
    ):
        super().__init__()
        # Attributes
        if size_h is None:
            size_h = size_in
        if size_out is None:
            size_out = size_in

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        if not legacy:
            self.bn_0 = CBatchNorm1d(c_dim, size_in, norm_method=norm_method)
            self.bn_1 = CBatchNorm1d(c_dim, size_h, norm_method=norm_method)
        else:
            self.bn_0 = CBatchNorm1d_legacy(c_dim, size_in, norm_method=norm_method)
            self.bn_1 = CBatchNorm1d_legacy(c_dim, size_h, norm_method=norm_method)

        self.fc_0 = nn.Conv1d(size_in, size_h, 1)
        self.fc_1 = nn.Conv1d(size_h, size_out, 1)
        self.actvn = nn.ReLU()

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = nn.Conv1d(size_in, size_out, 1, bias=False)
        # Initialization
        nn.init.zeros_(self.fc_1.weight)

    def forward(self, x, c):
        net = self.fc_0(self.actvn(self.bn_0(x, c)))
        dx = self.fc_1(self.actvn(self.bn_1(net, c)))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx


class ResnetBlockConv1d(nn.Module):
    """1D-Convolutional ResNet block class.

    Args:
        size_in (int): input dimension
        size_out (int): output dimension
        size_h (int): hidden dimension
    """

    def __init__(self, size_in, size_h=None, size_out=None):
        super().__init__()
        # Attributes
        if size_h is None:
            size_h = size_in
        if size_out is None:
            size_out = size_in

        self.size_in = size_in
        self.size_h = size_h
        self.size_out = size_out
        # Submodules
        self.bn_0 = nn.BatchNorm1d(size_in)
        self.bn_1 = nn.BatchNorm1d(size_h)

        self.fc_0 = nn.Conv1d(size_in, size_h, 1)
        self.fc_1 = nn.Conv1d(size_h, size_out, 1)
        self.actvn = nn.ReLU()

        if size_in == size_out:
            self.shortcut = None
        else:
            self.shortcut = nn.Conv1d(size_in, size_out, 1, bias=False)

        # Initialization
        nn.init.zeros_(self.fc_1.weight)

    def forward(self, x):
        net = self.fc_0(self.actvn(self.bn_0(x)))
        dx = self.fc_1(self.actvn(self.bn_1(net)))

        if self.shortcut is not None:
            x_s = self.shortcut(x)
        else:
            x_s = x

        return x_s + dx


# Utility modules
class AffineLayer(nn.Module):
    """Affine layer class.

    Args:
        c_dim (tensor): dimension of latent conditioned code c
        dim (int): input dimension
    """

    def __init__(self, c_dim, dim=3):
        super().__init__()
        self.c_dim = c_dim
        self.dim = dim
        # Submodules
        self.fc_A = nn.Linear(c_dim, dim * dim)
        self.fc_b = nn.Linear(c_dim, dim)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.fc_A.weight)
        nn.init.zeros_(self.fc_b.weight)
        with torch.no_grad():
            self.fc_A.bias.copy_(torch.eye(3).view(-1))
            self.fc_b.bias.copy_(torch.tensor([0.0, 0.0, 2.0]))

    def forward(self, x, p):
        assert x.size(0) == p.size(0)
        assert p.size(2) == self.dim
        batch_size = x.size(0)
        A = self.fc_A(x).view(batch_size, 3, 3)
        b = self.fc_b(x).view(batch_size, 1, 3)
        out = p @ A + b
        return out


class CBatchNorm1d(nn.Module):
    """Conditional batch normalization layer class.

    Args:
        c_dim (int): dimension of latent conditioned code c
        f_dim (int): feature dimension
        norm_method (str): normalization method
    """

    def __init__(self, c_dim, f_dim, norm_method="batch_norm"):
        super().__init__()
        self.c_dim = c_dim
        self.f_dim = f_dim
        self.norm_method = norm_method
        # Submodules
        self.conv_gamma = nn.Conv1d(c_dim, f_dim, 1)
        self.conv_beta = nn.Conv1d(c_dim, f_dim, 1)
        if norm_method == "batch_norm":
            self.bn = nn.BatchNorm1d(f_dim, affine=False)
        elif norm_method == "instance_norm":
            self.bn = nn.InstanceNorm1d(f_dim, affine=False)
        elif norm_method == "group_norm":
            self.bn = nn.GroupNorm1d(f_dim, affine=False)
        else:
            raise ValueError("Invalid normalization method!")
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.conv_gamma.weight)
        nn.init.zeros_(self.conv_beta.weight)
        nn.init.ones_(self.conv_gamma.bias)
        nn.init.zeros_(self.conv_beta.bias)

    def forward(self, x, c):
        assert x.size(0) == c.size(0)
        assert c.size(1) == self.c_dim

        # c is assumed to be of size batch_size x c_dim x T
        if len(c.size()) == 2:
            c = c.unsqueeze(2)

        # Affine mapping
        gamma = self.conv_gamma(c)
        beta = self.conv_beta(c)

        # Batchnorm
        net = self.bn(x)
        out = gamma * net + beta

        return out


class CBatchNorm1d_legacy(nn.Module):
    """Conditional batch normalization legacy layer class.

    Args:
        c_dim (int): dimension of latent conditioned code c
        f_dim (int): feature dimension
        norm_method (str): normalization method
    """

    def __init__(self, c_dim, f_dim, norm_method="batch_norm"):
        super().__init__()
        self.c_dim = c_dim
        self.f_dim = f_dim
        self.norm_method = norm_method
        # Submodules
        self.fc_gamma = nn.Linear(c_dim, f_dim)
        self.fc_beta = nn.Linear(c_dim, f_dim)
        if norm_method == "batch_norm":
            self.bn = nn.BatchNorm1d(f_dim, affine=False)
        elif norm_method == "instance_norm":
            self.bn = nn.InstanceNorm1d(f_dim, affine=False)
        elif norm_method == "group_norm":
            self.bn = nn.GroupNorm1d(f_dim, affine=False)
        else:
            raise ValueError("Invalid normalization method!")
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.zeros_(self.fc_gamma.weight)
        nn.init.zeros_(self.fc_beta.weight)
        nn.init.ones_(self.fc_gamma.bias)
        nn.init.zeros_(self.fc_beta.bias)

    def forward(self, x, c):
        batch_size = x.size(0)
        # Affine mapping
        gamma = self.fc_gamma(c)
        beta = self.fc_beta(c)
        gamma = gamma.view(batch_size, self.f_dim, 1)
        beta = beta.view(batch_size, self.f_dim, 1)
        # Batchnorm
        net = self.bn(x)
        out = gamma * net + beta

        return out

NCF_policies/algo2/ncf/ndf/vnn_occupancy_net_pointnet_dgcnn.py
carolinahiguera/NCF-Policy
from typing import Any
import pytorch_lightning as pl
import torch
import torch.nn as nn
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import cv2
import numpy as np

DEBUG = False


class NCF_Pipeline(pl.LightningModule):
    def __init__(self, cfg, digit_vae, ncf, ndf):
        super().__init__()

        self.cfg = cfg
        self.seq_len = cfg.seq_len
        self.digit_emb_dim = cfg.vae_params.latent_dim

        self.digit_vae = digit_vae
        self.ncf = ncf
        self.ndf = ndf

        self.mse_loss = nn.MSELoss(reduction="none")
        self.bce_loss = nn.BCELoss(reduction="none")
        self.l1_loss = nn.L1Loss(reduction="none")

    def configure_optimizers(self):
        # Using a scheduler is optional but can be helpful.
        # The scheduler reduces the LR if the validation performance hasn't improved for the last N epochs
        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        #     optimizer, mode="min", factor=0.2, patience=20, min_lr=5e-5
        # )

        lr = self.cfg.lr if self.cfg.ncf_arch == "transformer" else self.cfg.lr * 0.5

        optimizer = torch.optim.Adam(self.parameters(), lr=lr)

        if self.cfg.ncf_arch == "transformer":
            scheduler = torch.optim.lr_scheduler.StepLR(
                optimizer, step_size=5, gamma=0.1
            )
        else:
            scheduler = torch.optim.lr_scheduler.StepLR(
                optimizer, step_size=2, gamma=0.1
            )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
            },
        }
       
    def _get_digit_embeddings(self, images):
        batch_size = images.shape[0]
        seq_len = images.shape[1]
        images_hat = torch.zeros_like(images)

        digit_embeddings = torch.zeros(self.seq_len, batch_size, self.digit_emb_dim).to(
            images.device
        )
        for i in range(seq_len):
            s = images[:, i, :, :, :].unsqueeze(dim=1)
            self.digit_vae.eval()
            x_hat, emb = self.digit_vae(s)
            images_hat[:, i, :, :, :] = x_hat
            digit_embeddings[i, :, :] = emb

        return digit_embeddings.permute(1, 0, 2), images_hat

    def compute_loss_mse(self, gt, pred):
        loss = self.mse_loss(pred, gt)
        return loss.mean()

    def compute_loss_l1(self, gt, pred):
        loss = self.l1_loss(pred, gt)
        return loss.mean()

    def compute_loss_bce(self, gt, pred):
        pred = pred.squeeze()
        gt = gt.squeeze()
        loss = self.bce_loss(pred, gt)
        return loss.mean()

    def compute_loss_classic(self, gt, pred):
        label = gt.squeeze()
        # label = (label + 1) / 2.
        loss = (
            -1
            * (
                label * torch.log(pred + 1e-5)
                + (1 - label) * torch.log(1 - pred + 1e-5)
            ).mean()
        )
        return loss
  
    def compute_loss(self, pred, batch):
        gt = batch["p_contact_oe_t0"]

        if self.cfg.loss == "mse":
            loss = self.compute_loss_mse(gt, pred)
        elif self.cfg.loss == "l1":
            loss = self.compute_loss_l1(gt, pred)
        elif self.cfg.loss == "bce":
            loss = self.compute_loss_bce(gt, pred)
        elif self.cfg.loss == "classic":
            loss = self.compute_loss_classic(gt, pred)
        else:
            return NotImplementedError

        loss = loss * self.cfg.weight_loss
        return loss

    def forward(self, batch):
        ee_pose = batch["ee_pose"]
        embs_left = batch["digits_emb_left"]
        embs_right = batch["digits_emb_right"]
        ndf_point_cloud = batch["ndf_point_cloud"]
        ndf_query_point_cloud = batch["ndf_query_point_cloud"]
       
        # get ndf embeddings
        ndf_input = {}
        ndf_input["coords"] = ndf_query_point_cloud
        ndf_input["point_cloud"] = ndf_point_cloud
        out_ndf = self.ndf(ndf_input)

        inputs_ncf = {
            "digit_emb_left": embs_left,
            "digit_emb_right": embs_right,
            "ee_pose": ee_pose,
            "ee_pose_1": batch["ee_pose_1"],
            "ee_pose_2": batch["ee_pose_2"],
            "emb_ndf": out_ndf["features"],
        }

        pred_contact = self.ncf(inputs_ncf)

        # out = {
        #     "pred_contact": pred_contact,
        #     "images_hat_left": images_hat_left,
        #     "images_hat_right": images_hat_right,
        # }
        # return out

        return pred_contact

    def training_step(self, batch, batch_idx):
        out_ncf = self.forward(batch)
        pred_contact_oe = out_ncf["pred_contact"]
        loss = self.compute_loss(pred_contact_oe, batch)
        self.log(
            "train_loss", loss, on_step=True, on_epoch=True, prog_bar=False, logger=True
        )
        return {"loss": loss, "pred_contact_oe": pred_contact_oe}

    def validation_step(self, batch, batch_idx):
        out_ncf = self.forward(batch)
        pred_contact_oe = out_ncf["pred_contact"]
        loss = self.compute_loss(pred_contact_oe, batch)
        self.log(
            "val_loss", loss, on_step=True, on_epoch=True, prog_bar=False, logger=True
        )
        return {"loss": loss, "pred_contact_oe": pred_contact_oe}

NCF_policies/algo2/ncf/pipeline.py
carolinahiguera/NCF-Policy
# Copyright (c) Meta Platforms, Inc. and affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# Extension of Hora policy to tactile sensing with the DIGIT
# Authors: Haozhi Qi, Sudharshan Suresh
# --------------------------------------------------------


"""
    SE(3) pose utilities 
"""

# from isaacgym import gymapi
import numpy as np
import torch
from scipy.spatial.transform import Rotation as R
import copy 

def tf_to_xyzquat_numpy(pose: torch.Tensor) -> torch.Tensor:
    """
    convert 4 x 4 transformation matrices to [x, y, z, qx, qy, qz, qw]
    """
    pose = np.atleast_3d(pose)

    r = R.from_matrix(np.array(pose[:, 0:3, 0:3]))
    q = r.as_quat()  # qx, qy, qz, qw
    t = pose[:, :3, 3]
    xyz_quat = np.concatenate((t, q), axis=1)

    return xyz_quat  # (N, 7)


def xyzquat_to_tf_numpy(position_quat: np.ndarray) -> np.ndarray:
    """
    convert [x, y, z, qx, qy, qz, qw] to 4 x 4 transformation matrices
    """
    # try:
    position_quat = np.atleast_2d(position_quat)  # (N, 7)
    N = position_quat.shape[0]
    T = np.zeros((N, 4, 4))
    T[:, 0:3, 0:3] = R.from_quat(position_quat[:, 3:]).as_matrix()
    T[:, :3, 3] = position_quat[:, :3]
    T[:, 3, 3] = 1
    # except ValueError:
    #     print("Zero quat error!")
    return T.squeeze()


def tf2isaacgympose(cam_pose):
    r = R.from_matrix(cam_pose[:3, :3])
    axis_angle = r.as_rotvec()
    norm = np.linalg.norm(axis_angle) 
    transform = gymapi.Transform()
    transform.p = gymapi.Vec3(
        cam_pose[0, 3], cam_pose[1, 3], cam_pose[2, 3]
    )
    transform.r = gymapi.Quat.from_axis_angle(
        gymapi.Vec3(axis_angle[0], axis_angle[1], axis_angle[2]), norm
    )
    return transform

def flip_camera_z(cam_pose):
    '''
    flip camera pose by 180 deg in the X-axis, so depth is negative Z
    '''
    quat_real_to_tac_neural = gymapi.Quat.from_euler_zyx(np.radians(180), 0.0, 0.0)
    flip_cam_pose = copy.deepcopy(cam_pose)
    flip_cam_pose.r = cam_pose.r * quat_real_to_tac_neural
    return flip_cam_pose

def isaacgym2tacneural(cam_pose):
    cam_pose = np.array(
        [
            cam_pose.p.x,
            cam_pose.p.y,
            cam_pose.p.z,
            cam_pose.r.x,
            cam_pose.r.y,
            cam_pose.r.z,
            cam_pose.r.w,
        ]
    )
    return xyzquat_to_tf_numpy(cam_pose)
    
def xyzw_to_wxyz(quat: torch.Tensor) -> torch.Tensor:
    """
    Convention change: [x, y, z, qx, qy, qz, qw] --> [x, y, z, qw, qx, qy, qz]
    """
    if quat.shape[1] == 7:
        return quat[:, [0, 1, 2, 6, 3, 4, 5]]
    else:
        return quat[:, [3, 0, 1, 2]]


def wxyz_to_xyzw(quat: torch.Tensor) -> torch.Tensor:
    """
    Convention change: [x, y, z, qw, qx, qy, qz] --> [x, y, z, qx, qy, qz, qw]
    """
    if quat.shape[1] == 7:
        return quat[:, [0, 1, 2, 4, 5, 6, 3]]
    else:
        return quat[:, [1, 2, 3, 0]]


def transform_pc(pointclouds: np.ndarray, poses: np.ndarray):
    """
    Transform pointclouds by poses
    """

    if type(pointclouds) is not list:
        temp = pointclouds
        pointclouds = [None] * 1
        pointclouds[0] = temp
        poses = np.expand_dims(poses, axis=2)

    if len(poses.shape) < 3:
        poses = xyzquat_to_tf_numpy(poses)

    transformed_pointclouds = pointclouds
    # TODO: vectorize
    for i, (pointcloud, pose) in enumerate(zip(pointclouds, poses)):
        pointcloud = pointcloud.T
        # 3D affine transform
        pointcloud = pose @ np.vstack([pointcloud, np.ones((1, pointcloud.shape[1]))])
        pointcloud = pointcloud / pointcloud[3, :]
        pointcloud = pointcloud[:3, :].T
        transformed_pointclouds[i] = pointcloud
    return (
        transformed_pointclouds[0] if len(pointclouds) == 1 else transformed_pointclouds
    )


def wrap_angles(angles: torch.Tensor) -> torch.Tensor:
    """
    angles : (N, 3) angles in degrees
    Wraps to [-np.pi, np.pi] or [-180, 180]
    """

    mask = angles > 180.0
    angles[mask] -= 2.0 * 180.0

    mask = angles < -180.0
    angles[mask] += 2.0 * 180.0
    return angles


def quat2euler(quat: np.ndarray) -> np.ndarray:
    """
    Convert quaternions to euler angles
    """
    quat = np.atleast_2d(quat)
    r = R.from_quat(quat)
    return r.as_euler("xyz", degrees=True)


def rot2euler(rot: torch.Tensor) -> torch.Tensor:
    """
    Convert rotation matrix to euler angles
    Adapted from so3_rotation_angle() in  https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/so3.html
    """
    rot_trace = rot[:, 0, 0] + rot[:, 1, 1] + rot[:, 2, 2]
    phi_cos = torch.acos((rot_trace - 1.0) * 0.5)
    return torch.rad2deg(phi_cos)

    # r = R.from_matrix(np.atleast_3d(rot.cpu().numpy()))
    # eul = r.as_euler('xyz', degrees = True)
    # return torch.tensor(eul, device= rot.device)


def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
    """
    https://pytorch3d.readthedocs.io/en/latest/modules/transforms.html#pytorch3d.transforms.euler_angles_to_matrix
    Convert rotations given as Euler angles in radians to rotation matrices.
    Args:
        euler_angles: Euler angles in radians as tensor of shape (..., 3).
        convention: Convention string of three uppercase letters from
            {"X", "Y", and "Z"}.
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
        raise ValueError("Invalid input euler angles.")
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    matrices = [
        _axis_angle_rotation(c, e)
        for c, e in zip(convention, torch.unbind(euler_angles, -1))
    ]
    # return functools.reduce(torch.matmul, matrices)
    return matrices[0] @ matrices[1] @ matrices[2]


def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
    """
    https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Return the rotation matrices for one of the rotations about an axis
    of which Euler angles describe, for each value of the angle given.
    Args:
        axis: Axis label "X" or "Y or "Z".
        angle: any shape tensor of Euler angles in radians
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """

    cos = torch.cos(angle)
    sin = torch.sin(angle)
    one = torch.ones_like(angle)
    zero = torch.zeros_like(angle)

    if axis == "X":
        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
    elif axis == "Y":
        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
    elif axis == "Z":
        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
    else:
        raise ValueError("letter must be either X, Y or Z.")

    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))

def skew_matrix(v: np.ndarray) -> np.ndarray:
    """
    Get skew-symmetric matrix from vector
    """
    v = np.atleast_2d(v)
    # vector to its skew matrix
    mat = np.zeros((3, 3, v.shape[0]))
    mat[0, 1, :] = -1 * v[:, 2]
    mat[0, 2, :] = v[:, 1]

    mat[1, 0, :] = v[:, 2]
    mat[1, 2, :] = -1 * v[:, 0]

    mat[2, 0, :] = -1 * v[:, 1]
    mat[2, 1, :] = v[:, 0]
    return mat


def pose_from_vertex_normal(
    vertices: np.ndarray, normals: np.ndarray, shear_mag: float, delta: np.ndarray
) -> np.ndarray:
    """
    Generate SE(3) pose given
    vertices: (N, 3), normals: (N, 3), shear_mag: scalar, delta: (N, 1)
    """
    vertices = np.atleast_2d(vertices)
    normals = np.atleast_2d(normals)

    num_samples = vertices.shape[0]
    T = np.zeros((num_samples, 4, 4))  # transform from point coord to world coord
    T[:, 3, 3] = 1
    T[:, :3, 3] = vertices  # t

    # resolve ambiguous DoF
    """Find rotation of shear_vector so its orientation matches normal: np.dot(Rot, shear_vector) = normal
    https://math.stackexchange.com/questions/293116/rotating-one-3d-vector-to-another """

    cos_shear_mag = np.random.uniform(
        low=np.cos(shear_mag), high=1.0, size=(num_samples,)
    )  # Base of shear cone
    shear_phi = np.random.uniform(
        low=0.0, high=2 * np.pi, size=(num_samples,)
    )  # Circle of shear cone

    # Axis v = (shear_vector \cross normal)/(||shear_vector \cross normal||)
    shear_vector = np.array(
        [
            np.sqrt(1 - cos_shear_mag**2) * np.cos(shear_phi),
            np.sqrt(1 - cos_shear_mag**2) * np.sin(shear_phi),
            cos_shear_mag,
        ]
    ).T
    shear_vector_skew = skew_matrix(shear_vector)
    v = np.einsum("ijk,jk->ik", shear_vector_skew, normals.T).T
    v = v / np.linalg.norm(v, axis=1).reshape(-1, 1)

    # find corner cases
    check = np.einsum("ij,ij->i", normals, np.array([[0, 0, 1]]))
    zero_idx_up = check > 0.9  # pointing up
    zero_idx_down = check < -0.9  # pointing down

    v_skew, sampledNormals_skew = skew_matrix(v), skew_matrix(normals)

    # Angle theta = \arccos(z_axis \dot normal)
    # elementwise: theta = np.arccos(np.dot(shear_vector,normal)/(np.linalg.norm(shear_vector)*np.linalg.norm(normal)))
    theta = np.arccos(
        np.einsum("ij,ij->i", shear_vector, normals)
        / (np.linalg.norm(shear_vector, axis=1) * np.linalg.norm(normals, axis=1))
    )

    identity_3d = np.zeros(v_skew.shape)
    np.einsum("iij->ij", identity_3d)[:] = 1
    # elementwise: Rot = np.identity(3) + v_skew*np.sin(theta) + np.linalg.matrix_power(v_skew,2) * (1-np.cos(theta)) # rodrigues
    Rot = (
        identity_3d
        + v_skew * np.sin(theta)
        + np.einsum("ijn,jkn->ikn", v_skew, v_skew) * (1 - np.cos(theta))
    )  # rodrigues

    if np.any(zero_idx_up):
        Rot[:3, :3, zero_idx_up] = np.dstack([np.identity(3)] * np.sum(zero_idx_up))
    if np.any(zero_idx_down):
        Rot[:3, :3, zero_idx_down] = np.dstack(
            [np.array([[1, 0, 0], [0, -1, -0], [0, 0, -1]])] * np.sum(zero_idx_down)
        )

    # Rotation about Z axis is still ambiguous, generating random rotation b/w [0, 2pi] about normal axis
    # elementwise: RotDelta = np.identity(3) + normal_skew*np.sin(delta[i]) + np.linalg.matrix_power(normal_skew,2) * (1-np.cos(delta[i])) # rodrigues
    RotDelta = (
        identity_3d
        + sampledNormals_skew * np.sin(delta)
        + np.einsum("ijn,jkn->ikn", sampledNormals_skew, sampledNormals_skew)
        * (1 - np.cos(delta))
    )  # rodrigues

    # elementwise:  RotDelta @ Rot
    tfs = np.einsum("ijn,jkn->ikn", RotDelta, Rot)
    T[:, :3, :3] = np.rollaxis(tfs, 2)
    return T


def cam2gel(cam_pose, cam_dist):
    """
    Convert cam_pose to gel_pose
    """
    cam_tf = torch.eye(4, device=cam_pose.device)
    cam_tf[2, 3] = -cam_dist
    return cam_pose @ cam_tf[None, :]
NCF_policies/algo2/ncf/utils/pose.py
carolinahiguera/NCF-Policy
import torch
import numpy as np


def transform_pcd(pcd, transform):
    if pcd.shape[1] != 4:
        pcd = np.concatenate((pcd, np.ones((pcd.shape[0], 1))), axis=1)
    pcd_new = np.matmul(transform, pcd.T)[:-1, :].T
    return pcd_new

def transform_pcd_torch(pcd, transform):
    if pcd.dim() < 3:
        if pcd.shape[1] != 4:
            ones = torch.ones((pcd.shape[0], 1)).float().to(pcd.device)
            pcd = torch.cat((pcd, ones), dim=1)
        pcd_new = torch.matmul(transform, pcd.T)[:-1, :].T
    else:
        if pcd.shape[2] != 4:
            ones = torch.ones((pcd.shape[0], pcd.shape[1]))[:, :, None].float().to(pcd.device)
            pcd = torch.cat((pcd, ones), dim=2)
        pcd = pcd.transpose(1, 2)
        pcd_new = torch.matmul(transform, pcd)
        pcd_new = pcd_new.transpose(2, 1)[:, :, :-1]
    return pcd_new


EPS = 1e-8


def expmap2rotmat(r):
    """
    :param r: Axis-angle, Nx3
    :return: Rotation matrix, Nx3x3
    """
    dev = r.device
    assert r.shape[1] == 3
    bs = r.shape[0]
    theta = torch.sqrt(torch.sum(torch.pow(r, 2), 1, keepdim=True))
    cos_theta = torch.cos(theta).unsqueeze(-1)
    sin_theta = torch.sin(theta).unsqueeze(-1)
    eye = torch.unsqueeze(torch.eye(3), 0).repeat(bs, 1, 1).to(dev)
    norm_r = r / (theta + EPS)
    r_1 = torch.unsqueeze(norm_r, 2)  # N, 3, 1
    r_2 = torch.unsqueeze(norm_r, 1)  # N, 1, 3
    zero_col = torch.zeros(bs, 1).to(dev)
    skew_sym = torch.cat([zero_col, -norm_r[:, 2:3], norm_r[:, 1:2], norm_r[:, 2:3], zero_col,
                          -norm_r[:, 0:1], -norm_r[:, 1:2], norm_r[:, 0:1], zero_col], 1)
    skew_sym = skew_sym.contiguous().view(bs, 3, 3)
    R = cos_theta*eye + (1-cos_theta)*torch.bmm(r_1, r_2) + sin_theta*skew_sym
    return R


def rotmat2expmap(R):
    """
    :param R: Rotation matrix, Nx3x3
    :return: r: Rotation vector, Nx3
    """
    assert R.shape[1] == R.shape[2] == 3
    theta = torch.acos(torch.clamp((R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2] - 1) / 2, min=-1., max=1.)).view(-1, 1)
    r = torch.stack((R[:, 2, 1]-R[:, 1, 2], R[:, 0, 2]-R[:, 2, 0], R[:, 1, 0]-R[:, 0, 1]), 1) / (2*torch.sin(theta))
    r_norm = r / torch.sqrt(torch.sum(torch.pow(r, 2), 1, keepdim=True))
    return theta * r_norm


def quat2expmap(q):
    """
    :param q: quaternion, Nx4
    :return: r: Axis-angle, Nx3
    """
    assert q.shape[1] == 4
    cos_theta_2 = torch.clamp(q[:, 0: 1], min=-1., max=1.)
    theta = torch.acos(cos_theta_2)*2
    sin_theta_2 = torch.sqrt(1-torch.pow(cos_theta_2, 2))
    r = theta * q[:, 1:4] / (sin_theta_2 + EPS)
    return r


def expmap2quat(r):
    """
    :param r: Axis-angle, Nx3
    :return: q: quaternion, Nx4
        """
    assert r.shape[1] == 3
    theta = torch.sqrt(torch.sum(torch.pow(r, 2), 1, keepdim=True))
    unit_r = r / theta
    theta_2 = theta / 2.
    cos_theta_2 = torch.cos(theta_2)
    sin_theta_2 = torch.sin(theta_2)
    q = torch.cat((cos_theta_2, unit_r*sin_theta_2), 1)
    return q


def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor:
    """Convert quaternion vector to angle axis of rotation.

    Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h

    Args:
        quaternion (torch.Tensor): tensor with quaternions.

    Return:
        torch.Tensor: tensor with angle axis of rotation.

    Shape:
        - Input: :math:`(*, 4)` where `*` means, any number of dimensions
        - Output: :math:`(*, 3)`

    Example:
        >>> quaternion = torch.rand(2, 4)  # Nx4
        >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion)  # Nx3
    """
    if not torch.is_tensor(quaternion):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(quaternion)))

    if not quaternion.shape[-1] == 4:
        raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}"
                         .format(quaternion.shape))
    # unpack input and compute conversion
    q1: torch.Tensor = quaternion[..., 1]
    q2: torch.Tensor = quaternion[..., 2]
    q3: torch.Tensor = quaternion[..., 3]
    sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3

    sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta)
    cos_theta: torch.Tensor = quaternion[..., 0]
    two_theta: torch.Tensor = 2.0 * torch.where(
        cos_theta < 0.0,
        torch.atan2(-sin_theta, -cos_theta),
        torch.atan2(sin_theta, cos_theta))

    k_pos: torch.Tensor = two_theta / sin_theta
    k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta)
    k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg)

    angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3]
    angle_axis[..., 0] += q1 * k
    angle_axis[..., 1] += q2 * k
    angle_axis[..., 2] += q3 * k
    return angle_axis

def angle_axis_to_rotation_matrix(angle_axis):
    """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix

    Args:
        angle_axis (Tensor): tensor of 3d vector of axis-angle rotations.

    Returns:
        Tensor: tensor of 4x4 rotation matrices.

    Shape:
        - Input: :math:`(N, 3)`
        - Output: :math:`(N, 4, 4)`

    Example:
        >>> input = torch.rand(1, 3)  # Nx3
        >>> output = tgm.angle_axis_to_rotation_matrix(input)  # Nx4x4
    """
    def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6):
        # We want to be careful to only evaluate the square root if the
        # norm of the angle_axis vector is greater than zero. Otherwise
        # we get a division by zero.
        k_one = 1.0
        theta = torch.sqrt(theta2)
        wxyz = angle_axis / (theta + eps)
        wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
        cos_theta = torch.cos(theta)
        sin_theta = torch.sin(theta)

        r00 = cos_theta + wx * wx * (k_one - cos_theta)
        r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
        r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
        r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
        r11 = cos_theta + wy * wy * (k_one - cos_theta)
        r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
        r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
        r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
        r22 = cos_theta + wz * wz * (k_one - cos_theta)
        rotation_matrix = torch.cat(
            [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
        return rotation_matrix.view(-1, 3, 3)

    def _compute_rotation_matrix_taylor(angle_axis):
        rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
        k_one = torch.ones_like(rx)
        rotation_matrix = torch.cat(
            [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1)
        return rotation_matrix.view(-1, 3, 3)

    # stolen from ceres/rotation.h

    _angle_axis = torch.unsqueeze(angle_axis, dim=1)
    theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
    theta2 = torch.squeeze(theta2, dim=1)

    # compute rotation matrices
    rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
    rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)

    # create mask to handle both cases
    eps = 1e-6
    mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
    mask_pos = (mask).type_as(theta2)
    mask_neg = (mask == False).type_as(theta2)  # noqa

    # create output pose matrix
    batch_size = angle_axis.shape[0]
    rotation_matrix = torch.eye(4).to(angle_axis.device).type_as(angle_axis)
    rotation_matrix = rotation_matrix.view(1, 4, 4).repeat(batch_size, 1, 1)
    # fill output matrix with masked values
    rotation_matrix[..., :3, :3] = \
        mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
    return rotation_matrix  # Nx4x4        


def rotation_matrix_to_angle_axis(rotation_matrix):
    """Convert 3x4 rotation matrix to Rodrigues vector

    Args:
        rotation_matrix (Tensor): rotation matrix.

    Returns:
        Tensor: Rodrigues vector transformation.

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 3)`

    Example:
        >>> input = torch.rand(2, 3, 4)  # Nx4x4
        >>> output = tgm.rotation_matrix_to_angle_axis(input)  # Nx3
    """
    # todo add check that matrix is a valid rotation matrix
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)
    return quaternion_to_angle_axis(quaternion)


def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6):
    """Convert 3x4 rotation matrix to 4d quaternion vector

    This algorithm is based on algorithm described in
    https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201

    Args:
        rotation_matrix (Tensor): the rotation matrix to convert.

    Return:
        Tensor: the rotation in quaternion

    Shape:
        - Input: :math:`(N, 3, 4)`
        - Output: :math:`(N, 4)`

    Example:
        >>> input = torch.rand(4, 3, 4)  # Nx3x4
        >>> output = tgm.rotation_matrix_to_quaternion(input)  # Nx4
    """
    if not torch.is_tensor(rotation_matrix):
        raise TypeError("Input type is not a torch.Tensor. Got {}".format(
            type(rotation_matrix)))

    if len(rotation_matrix.shape) > 3:
        raise ValueError(
            "Input size must be a three dimensional tensor. Got {}".format(
                rotation_matrix.shape))
    if not rotation_matrix.shape[-2:] == (3, 4):
        raise ValueError(
            "Input size must be a N x 3 x 4  tensor. Got {}".format(
                rotation_matrix.shape))

    rmat_t = torch.transpose(rotation_matrix, 1, 2)

    mask_d2 = rmat_t[:, 2, 2] < eps

    mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1]
    mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1]

    t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1)
    t0_rep = t0.repeat(4, 1).t()

    t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2]
    q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] + rmat_t[:, 1, 0],
                      t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1)
    t1_rep = t1.repeat(4, 1).t()

    t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0],
                      rmat_t[:, 2, 0] + rmat_t[:, 0, 2],
                      rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1)
    t2_rep = t2.repeat(4, 1).t()

    t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2]
    q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1],
                      rmat_t[:, 2, 0] - rmat_t[:, 0, 2],
                      rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1)
    t3_rep = t3.repeat(4, 1).t()

    mask_c0 = mask_d2 * mask_d0_d1
    mask_c1 = mask_d2 * (1 - mask_d0_d1)
    mask_c2 = (1 - mask_d2) * mask_d0_nd1
    mask_c3 = (1 - mask_d2) * (1 - mask_d0_nd1)
    mask_c0 = mask_c0.view(-1, 1).type_as(q0)
    mask_c1 = mask_c1.view(-1, 1).type_as(q1)
    mask_c2 = mask_c2.view(-1, 1).type_as(q2)
    mask_c3 = mask_c3.view(-1, 1).type_as(q3)

    q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3
    q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 +  # noqa
                    t2_rep * mask_c2 + t3_rep * mask_c3)  # noqa
    q *= 0.5
    return q



NCF_policies/algo2/ncf/utils/torch_util.py
carolinahiguera/NCF-Policy
import torch
import numpy as np
from scipy.spatial.transform import Rotation as R

from algo2.ncf.utils import torch_util

EEF_MAX = np.array([0.10, 0.10, 0.20])  # m
DELTA_EEF_MAX = np.array([1e-2, 1e-2, 2e-2])  # cm
DELTA_QUAT_MAX = np.array([0.01, 0.01, 0.01, 1.0])


class NCF_dataloader:
    def __init__(
        self,
        num_envs,
        path_mesh_object,
        path_pointcloud_object,
        path_ndf_code_object,
        pc_subsample=1.0,
    ):
        self.num_envs = num_envs
        self.path_mesh_object = path_mesh_object
        self.path_pointcloud_object = path_pointcloud_object
        self.pc_subsample = pc_subsample

        self.shape_code = np.load(path_ndf_code_object)
        self.shape_code = torch.tensor(self.shape_code, dtype=torch.float32)

        self.ncf_point_cloud = np.load(path_pointcloud_object)
        self.ncf_point_cloud = self.ncf_point_cloud - np.mean(
            self.ncf_point_cloud, axis=0
        )
        self.n_points = len(self.ncf_point_cloud)


    def get_ee_sequence(self, ee_poses):
        num_envs = ee_poses.shape[0]
        ee_poses = ee_poses.cpu().numpy()
        # ee_poses[:, :, 2] -= 0.040

        ee_poses_envs = np.zeros_like(ee_poses)
        ee_t1_envs = np.zeros((num_envs, 7))
        ee_t2_envs = np.zeros((num_envs, 7))

        for i in range(num_envs):
            ee_seq = ee_poses[i]
            ee_seq_cp = np.copy(ee_seq)

            r0 = R.from_quat(ee_seq_cp[0][3:]).as_matrix()
            r1 = R.from_quat(ee_seq_cp[1][3:]).as_matrix()
            r2 = R.from_quat(ee_seq_cp[2][3:]).as_matrix()
            q_t_1 = R.from_matrix(np.matmul(r0, r1.T)).as_quat()
            q_t_2 = R.from_matrix(np.matmul(r0, r2.T)).as_quat()
            t_t_1 = ee_seq_cp[0][0:3] - ee_seq_cp[1][0:3]
            t_t_2 = ee_seq_cp[0][0:3] - ee_seq_cp[2][0:3]
            ee_t1 = np.concatenate((t_t_1, q_t_1))
            ee_t2 = np.concatenate((t_t_2, q_t_2))

            ee_poses_envs[i][:, 0:3] = ee_seq[:, 0:3] / EEF_MAX
            ee_poses_envs[i][:, 3:] = ee_seq[:, 3:]
            ee_t1_envs[i][0:3] = ee_t1[0:3] / DELTA_EEF_MAX
            ee_t1_envs[i][3:] = ee_t1[3:] / DELTA_QUAT_MAX
            ee_t2_envs[i][0:3] = ee_t2[0:3] / DELTA_EEF_MAX
            ee_t2_envs[i][3:] = ee_t2[3:] / DELTA_QUAT_MAX

        ee_poses_envs = torch.tensor(ee_poses_envs, dtype=torch.float32)
        ee_t1_envs = torch.tensor(ee_t1_envs, dtype=torch.float32)
        ee_t2_envs = torch.tensor(ee_t2_envs, dtype=torch.float32)

        return ee_poses_envs, ee_t1_envs, ee_t2_envs

    def get_pointcloud_data(self):
        idx_points = np.arange(self.n_points)
        n_query_pts = int(self.n_points * self.pc_subsample)
        idx_query = np.random.choice(idx_points, size=n_query_pts)
        ncf_query_points = self.ncf_point_cloud[idx_query]
        ncf_query_points = torch.tensor(ncf_query_points, dtype=torch.float32).repeat(
            self.num_envs, 1, 1
        )
        shape_code = self.shape_code.repeat(self.num_envs, 1)

        return shape_code, ncf_query_points, idx_query

    def get_pointcloud_data_ndf(self):
        # points for ncf

        ncf_point_cloud = np.load(self.path_pointcloud_object)
        ncf_point_cloud = ncf_point_cloud - np.mean(ncf_point_cloud, axis=0)
        n_points = len(ncf_point_cloud)

        idx_points = np.arange(n_points)
        n_query_pts = int(n_points * 1.0)
        idx_query = np.random.choice(idx_points, size=n_query_pts)
        ncf_query_points = ncf_point_cloud[idx_query]

        ncf_query_points = torch.tensor(ncf_query_points, dtype=torch.float32).repeat(
            self.num_envs, 1, 1
        )

        # points for ndf
        object_trimesh = trimesh.load(self.path_mesh_object)
        T = np.eye(4)
        T[0:3, 0:3] = R.from_euler("xyz", [0, 0, 90], degrees=True).as_matrix()
        object_trimesh = object_trimesh.apply_transform(T)

        ndf_point_cloud = np.array(
            trimesh.sample.sample_surface(object_trimesh, n_points // 2)[0]
        )
        ndf_point_cloud = ndf_point_cloud - np.mean(ndf_point_cloud, axis=0)
        ndf_point_cloud = torch.tensor(ndf_point_cloud, dtype=torch.float32).repeat(
            self.num_envs, 1, 1
        )

        return ndf_point_cloud, ncf_query_points, idx_query


import trimesh
from vedo import trimesh2vedo, Points, show, Text2D
from .pose import xyzquat_to_tf_numpy, euler_angles_to_matrix


class NCF_Viz:
    def __init__(self, path_assets=None, cupholder_scale=0.75):
        # path_assets = (
        #     "/home/chiguera/Documents/NCF/NCFgym/assets/factory/mesh/ncf_mug_cupholder/"
        # )
        if path_assets is None:
            path_assets = "/home/chiguera/Documents/NCF/NCFgym/assets/factory/mesh/ncf_mug_cupholder/"

        object_path = f"{path_assets}/obj_assets/1/obj_collision.obj"
        object_pc_path = f"{path_assets}/obj_assets/1/obj_pointcloud.npy"
        cupholder_path = f"{path_assets}/cupholder/cupholder_v2.obj"

        self.mesh_object = trimesh.load_mesh(object_path)
        self.mesh_object.apply_scale(1.0)
        T = np.eye(4)
        T[0:3, 0:3] = R.from_euler("xyz", [0.0, 0.0, 90.0], degrees=True).as_matrix()
        self.mesh_object = self.mesh_object.apply_transform(T)

        self.cupholder_trimesh = trimesh.load(cupholder_path)
        self.cupholder_trimesh = self.cupholder_trimesh.apply_scale(cupholder_scale)

        self.pc = np.load(object_pc_path)

        self.object_3d_cam = dict(
            position=(-7.54483e-3, -0.0849045, -0.250212),
            focal_point=(-4.82255e-3, -2.87705e-3, 0),
            viewup=(0.580505, -0.775588, 0.247946),
            distance=0.263329,
            clipping_range=(0.174192, 0.376170),
        )

    def show_obj_probabilities(self, dist, interactive=False):
        pc = Points(self.pc, r=15)
        pc = pc.cmap(
            "plasma",
            dist,
            vmin=0.0,
            vmax=1.0,
        )
        mesh_vedo = trimesh2vedo(self.mesh_object).clone()
        # mesh_vedo.subdivide(n=10, method=2)
        mesh_vedo = mesh_vedo.interpolate_data_from(
            pc, n=5, on="points", kernel="gaussian"
        ).cmap("plasma", vmin=0.0, vmax=1.0)

        if interactive:
            show([mesh_vedo, pc], axes=1, camera=self.object_3d_cam)
        else:
            img = show(
                [mesh_vedo, pc],
                axes=1,
                camera=self.object_3d_cam,
                interactive=False,
                offscreen=True,
            ).screenshot(asarray=True)
            return img

    def show_scene(self, object_pose, cupholder_pose, probs, label, interactive=False):
        # cam = dict(
        #     position=(0.227997, -0.375556, -0.239064),
        #     focal_point=(-0.0186972, 0.0437448, 0.0995462),
        #     viewup=(0.265986, -0.506072, 0.820452),
        #     distance=0.592729,
        #     clipping_range=(0.297838, 0.853488),
        # )
        cam = dict(
            position=(0.335978, -0.270905, -0.0235918),
            focal_point=(-0.0186971, 0.0437446, 0.0995460),
            viewup=(0.194292, -0.159759, 0.967847),
            distance=0.489859,
            clipping_range=(0.240417, 0.708860),
        )
        tf = np.eye(4)
        tf[0:3, 0:3] = euler_angles_to_matrix(
            euler_angles=torch.tensor([[0, 0, -np.pi]]), convention="XYZ"
        ).numpy()

        object_pose = xyzquat_to_tf_numpy(object_pose.cpu().numpy())  # @ tf
        cupholder_pose = xyzquat_to_tf_numpy(cupholder_pose.cpu().numpy())  # @ tf

        cupholder_i = self.cupholder_trimesh.copy()
        cupholder_i.apply_transform(cupholder_pose)

        obj_i = self.mesh_object.copy()
        obj_i.apply_transform(object_pose)

        object_pc_i = trimesh.points.PointCloud(self.pc.copy())
        object_pc_i.apply_transform(object_pose)

        mesh = trimesh2vedo(obj_i).clone()
        mesh.subdivide(n=3, method=2)

        holder_mesh = trimesh2vedo(cupholder_i).clone()

        pts_external_contact = Points(object_pc_i.vertices, r=10).cmap(
            "plasma", probs, vmin=0.0, vmax=1.0
        )

        mesh.interpolate_data_from(
            pts_external_contact, n=3, on="points", kernel="gaussian"
        ).cmap("plasma", vmin=0.0, vmax=1.0)

        holder_mesh.alpha(0.5).color("gray")

        if interactive:
            show([mesh, holder_mesh, pts_external_contact], axes=0, camera=cam).close()
        else:
            text = Text2D(label, pos="bottom-left", s=1, c="black")
            img = show(
                [mesh, holder_mesh, pts_external_contact, text],
                axes=0,
                camera=cam,
                interactive=False,
                offscreen=True,
            ).screenshot(asarray=True)
            show().close()
            return img

    def show_gt_pred(self, object_pose, cupholder_pose, p_gt, p_pred):
        img_gt = self.show_scene(
            object_pose, cupholder_pose, p_gt, "Ground Truth", interactive=False
        )
        img_pred = self.show_scene(
            object_pose, cupholder_pose, p_pred, "NCF estimation", interactive=False
        )
        return img_gt, img_pred

    def show_gt_sim_real(
        self, object_pose, cupholder_pose, p_gt, p_sim, p_sim_real, p_real
    ):
        img_gt = self.show_scene(
            object_pose, cupholder_pose, p_gt, "Ground Truth", interactive=False
        )
        img_sim = self.show_scene(
            object_pose, cupholder_pose, p_sim, "NCF sim", interactive=False
        )

        img_sim_real = self.show_scene(
            object_pose,
            cupholder_pose,
            p_sim_real,
            "NCF sim and real inputs",
            interactive=False,
        )

        img_real = self.show_scene(
            object_pose, cupholder_pose, p_real, "NCF real", interactive=False
        )
        return img_gt, img_sim, img_sim_real, img_real

NCF_policies/algo2/ncf/utils/utils_ncf.py
carolinahiguera/NCF-Policy
import torch
from torch import nn
from torch.nn import functional as F


class Interpolate(nn.Module):
    """nn.Module wrapper for F.interpolate."""

    def __init__(self, size=None, scale_factor=None):
        super().__init__()
        self.size, self.scale_factor = size, scale_factor

    def forward(self, x):
        return F.interpolate(x, size=self.size, scale_factor=self.scale_factor)


def resize_conv3x3(in_planes, out_planes, scale=1):
    """upsample + 3x3 convolution with padding to avoid checkerboard artifact."""
    if scale == 1:
        return conv3x3(in_planes, out_planes)
    return nn.Sequential(
        Interpolate(scale_factor=scale), conv3x3(in_planes, out_planes)
    )


def resize_conv1x1(in_planes, out_planes, scale=1):
    """upsample + 1x1 convolution with padding to avoid checkerboard artifact."""
    if scale == 1:
        return conv1x1(in_planes, out_planes)
    return nn.Sequential(
        Interpolate(scale_factor=scale), conv1x1(in_planes, out_planes)
    )


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution."""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding."""
    return nn.Conv2d(
        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
    )


class EncoderBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super().__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class DecoderBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, scale=1, upsample=None):
        super().__init__()
        self.conv1 = resize_conv3x3(inplanes, inplanes)
        self.bn1 = nn.BatchNorm2d(inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = resize_conv3x3(inplanes, planes, scale)
        self.bn2 = nn.BatchNorm2d(planes)
        self.upsample = upsample

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.upsample is not None:
            identity = self.upsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNetEncoder(nn.Module):
    def __init__(self, block, layers, first_conv=False, maxpool1=False):
        super().__init__()

        self.inplanes = 64
        self.first_conv = first_conv
        self.maxpool1 = maxpool1

        if self.first_conv:
            self.conv1 = nn.Conv2d(
                3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
            )
        else:
            self.conv1 = nn.Conv2d(
                3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False
            )

        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)

        if self.maxpool1:
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        else:
            self.maxpool = nn.MaxPool2d(kernel_size=1, stride=1)

        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x


class ResNetDecoder(nn.Module):
    """Resnet in reverse order."""

    def __init__(
        self, block, layers, latent_dim, input_height, first_conv=False, maxpool1=False
    ):
        super().__init__()

        self.expansion = block.expansion
        self.inplanes = 512 * block.expansion
        self.first_conv = first_conv
        self.maxpool1 = maxpool1
        self.input_height = input_height

        self.upscale_factor = 8

        self.linear = nn.Linear(latent_dim, self.inplanes * 4 * 4)

        self.layer1 = self._make_layer(block, 256, layers[0], scale=2)
        self.layer2 = self._make_layer(block, 128, layers[1], scale=2)
        self.layer3 = self._make_layer(block, 64, layers[2], scale=2)

        if self.maxpool1:
            self.layer4 = self._make_layer(block, 64, layers[3], scale=2)
            self.upscale_factor *= 2
        else:
            self.layer4 = self._make_layer(block, 64, layers[3])

        if self.first_conv:
            self.upscale = Interpolate(scale_factor=2)
            self.upscale_factor *= 2
        else:
            self.upscale = Interpolate(scale_factor=1)

        # interpolate after linear layer using scale factor
        self.upscale1 = Interpolate(size=input_height // self.upscale_factor)

        self.conv1 = nn.Conv2d(
            64 * block.expansion, 3, kernel_size=3, stride=1, padding=1, bias=False
        )

    def _make_layer(self, block, planes, blocks, scale=1):
        upsample = None
        if scale != 1 or self.inplanes != planes * block.expansion:
            upsample = nn.Sequential(
                resize_conv1x1(self.inplanes, planes * block.expansion, scale),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, scale, upsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.linear(x)

        # NOTE: replaced this by Linear(in_channels, 514 * 4 * 4)
        # x = F.interpolate(x, scale_factor=4)

        x = x.view(x.size(0), 512 * self.expansion, 4, 4)
        x = self.upscale1(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.upscale(x)

        x = self.conv1(x)
        return x


def resnet18_encoder(first_conv, maxpool1):
    return ResNetEncoder(EncoderBlock, [2, 2, 2, 2], first_conv, maxpool1)


def resnet18_decoder(latent_dim, input_height, first_conv, maxpool1):
    return ResNetDecoder(
        DecoderBlock, [2, 2, 2, 2], latent_dim, input_height, first_conv, maxpool1
    )

NCF_policies/algo2/ncf/vae/resnet_modules.py
carolinahiguera/NCF-Policy
import pytorch_lightning as pl
from torch import nn
import torch
from torchvision.utils import make_grid

import numpy as np
from matplotlib.pyplot import imshow, figure

from .resnet_modules import resnet18_decoder, resnet18_encoder


class VAE(pl.LightningModule):
    def __init__(self, enc_out_dim=512, latent_dim=256, input_height=32):
        super().__init__()

        self.save_hyperparameters()

        # encoder, decoder
        self.encoder = resnet18_encoder(first_conv=False, maxpool1=False)
        self.decoder = resnet18_decoder(
            latent_dim=latent_dim,
            input_height=input_height,
            first_conv=False,
            maxpool1=False,
        )

        # distribution parameters
        self.fc_mu = nn.Linear(enc_out_dim, latent_dim)
        self.fc_var = nn.Linear(enc_out_dim, latent_dim)

        # for the gaussian likelihood
        self.log_scale = nn.Parameter(torch.Tensor([0.0]))

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-4)

    def gaussian_likelihood(self, x_hat, logscale, x):
        scale = torch.exp(logscale)
        mean = x_hat
        dist = torch.distributions.Normal(mean, scale)

        # measure prob of seeing image under p(x|z)
        log_pxz = dist.log_prob(x)
        return log_pxz.sum(dim=(1, 2, 3))

    def kl_divergence(self, z, mu, std):
        # --------------------------
        # Monte carlo KL divergence
        # --------------------------
        # 1. define the first two probabilities (in this case Normal for both)
        p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
        q = torch.distributions.Normal(mu, std)

        # 2. get the probabilities from the equation
        log_qzx = q.log_prob(z)
        log_pz = p.log_prob(z)

        # kl
        kl = log_qzx - log_pz
        kl = kl.sum(-1)
        return kl

    def training_step(self, batch, batch_idx):
        x = batch.permute(1, 0, 2, 3, 4)
        x = x.reshape(-1, *x.shape[2:])
        # x = x.permute(0,3,1,2)

        # encode x to get the mu and variance parameters
        x_encoded = self.encoder(x)
        mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded)

        # sample z from q
        std = torch.exp(log_var / 2)
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()

        # decoded
        x_hat = self.decoder(z)

        # reconstruction loss
        recon_loss = self.gaussian_likelihood(x_hat, self.log_scale, x)

        # kl
        kl = self.kl_divergence(z, mu, std)

        # elbo
        elbo = kl - recon_loss
        elbo = elbo.mean()

        self.log_dict(
            {
                "elbo": elbo,
                "kl": kl.mean(),
                "recon_loss": recon_loss.mean(),
                "reconstruction": recon_loss.mean(),
                "kl": kl.mean(),
            }
        )

        return {"loss": elbo, "pred": x_hat}

    def forward(self, batch):
        x = batch.permute(1, 0, 2, 3, 4)
        x = x.reshape(-1, *x.shape[2:])
        # x = x.permute(0,3,1,2)

        # encode x to get the mu and variance parameters
        x_encoded = self.encoder(x)
        mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded)

        # sample z from q
        std = torch.exp(log_var / 2)
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()

        # decoded
        x_hat = self.decoder(z)

        return x_hat, z

    def get_embedding(self, x):
        x_encoded = self.encoder(x)
        mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded)
        # sample z from q
        std = torch.exp(log_var / 2)
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()
        return z

    def validation_step(self, batch, batch_idx):
        x = batch.permute(1, 0, 2, 3, 4)
        x = x.reshape(-1, *x.shape[2:])
        z = self.get_embedding(x)
        x_hat = self.decoder(z)
        # reconstruction loss
        recon_loss = self.gaussian_likelihood(x_hat, self.log_scale, x)
        return {"loss": recon_loss, "pred": x_hat}


class VAE_Callback(pl.Callback):
    def __init__(self, every_n_epochs):
        super().__init__()
        self.img_size = None
        self.num_preds = 8
        self.every_n_epochs = every_n_epochs

    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0
    ):
        if trainer.current_epoch % self.every_n_epochs == 0:
            x = batch.permute(1, 0, 2, 3, 4)
            x = x.reshape(-1, *x.shape[2:])
            # idx = np.random.randint(0, x.size(0)-1, self.num_preds, )
            max_len = self.num_preds if self.num_preds <= x.size(0) else x.size(0)
            idx = np.arange(max_len)
            x_org = x[idx].detach().cpu()
            x_pred = outputs["pred"][idx].detach().cpu()
            # UNDO DATA NORMALIZATION
            # mean, std = np.array(0.5), np.array(0.5)
            img_pred = make_grid(x_pred).permute(1, 2, 0).numpy()  # * std + mean
            img_org = make_grid(x_org).permute(1, 2, 0).numpy()  # * std + mean
            img = np.vstack((img_org, img_pred))
            # PLOT IMAGES
            trainer.logger.experiment.add_image(
                f"reconstruction",
                torch.tensor(img).permute(2, 0, 1),
                global_step=trainer.global_step,
            )

    # def on_train_epoch_end(self, trainer, pl_module):
    #     if trainer.current_epoch % self.every_n_epochs == 0:
    #         rand_v = torch.rand((self.num_preds, pl_module.hparams.latent_dim), device=pl_module.device)
    #         p = torch.distributions.Normal(torch.zeros_like(rand_v), torch.ones_like(rand_v))
    #         z = p.rsample()

    #         # SAMPLE IMAGES
    #         with torch.no_grad():
    #             pred = pl_module.decoder(z.to(pl_module.device)).cpu()

    #         # UNDO DATA NORMALIZATION
    #         mean, std = np.array(0.5), np.array(0.5)
    #         img = make_grid(pred).permute(1, 2, 0).numpy() * std + mean

    #         # PLOT IMAGES
    #         trainer.logger.experiment.add_image(f'img_{trainer.current_epoch}',torch.tensor(img).permute(2, 0, 1), global_step=trainer.global_step)

NCF_policies/algo2/ncf/vae/vae.py
carolinahiguera/NCF-Policy
import os
import time
from matplotlib import pyplot as plt
import torch
import trimesh
import numpy as np
import cv2
from termcolor import cprint
from dataclasses import dataclass
from scipy.spatial.transform import Rotation as Rot

from vedo import Plotter, Mesh, Points, show, Text2D
from vedo import trimesh2vedo

from algo2.policy.deploy.franka import FrankaEnv
from algo2.policy.models.models import ActorCritic
from algo2.policy.models.running_mean_std import RunningMeanStd
from algo2.utils.misc import add_to_fifo
from algo2.policy.deploy.ncf_viz import NCFViz

from pynput import keyboard

EEF_MAX = np.array([0.10, 0.10, 0.20])  # m
DELTA_EEF_MAX = np.array([1e-2, 1e-2, 2e-2])  # cm
DELTA_QUAT_MAX = np.array([0.01, 0.01, 0.01, 1.0])


@dataclass
class Digit_VAE_Params:
    root_path: str = os.path.abspath(os.path.join(".."))
    image_size: int = 64
    channels: int = 3
    seq_len: int = 5
    latent_dim: int = 64
    enc_out_dim: int = 512
    source_data: str = "sim"
    checkpoint_dir = os.path.join(root_path, "digit_vae", "checkpoints")


@dataclass
class NCF_Params:
    root_path: str = os.path.abspath(os.path.join("."))
    seq_len: int = 10
    source_data: str = "sim"
    hidden_size_ncf: int = 128
    pc_subsample: float = 1.0
    vae_params = Digit_VAE_Params()


class HardwarePlayer(object):
    def __init__(self, config):
        current_file = os.path.abspath(__file__)
        self.root_dir = os.path.abspath(
            os.path.join(current_file, "..", "..", "..", "..", "..")
        )

        self.action_scale_pos = 0.2  # 0.2  # proprio only 0.3
        self.action_scale_rot = 0.3  # 0.45  # proprio only 0.2
        self.control_freq = 30
        self.device = torch.device("cuda:0")

        self.full_config = config
        self.network_config = config.train.network
        self.ppo_config = config.train.ppo
        self.task_config = config.task

        # ---- Env Info ----
        self.actions_num = 6
        self.obs_shape = (21,)

        # ---- Priv Info ----
        self.priv_info_dim = self.ppo_config["ncf_output_dim"]
        self.priv_info = self.ppo_config["ncf_info"]
        self.ncf_proprio_adapt = False

        # ---- Tactile Info ----
        self.tactile_info = self.ppo_config["tactile_info"]
        self.tactile_seq_length = self.ppo_config["tactile_seq_length"]
        self.tactile_info_embed_dim = self.ppo_config["tactile_info_embed_dim"]
        self.tactile_info_dim = (
            self.tactile_info_embed_dim * self.tactile_seq_length * 2
        )

        # ---- NCF Info ----
        self.ncf_info = self.ppo_config["ncf_info"]
        self.ncf_use_gt = self.ppo_config["ncf_use_gt"]
        self.ncf_output_dim = self.ppo_config["ncf_output_dim"]
        if self.ncf_info:
            self.tactile_info = False

        # ---- Model ----
        net_config = {
            "actor_units": self.network_config.mlp.units,
            "actions_num": self.actions_num,
            "proprio_input_shape": self.obs_shape,
            "tactile_info": self.tactile_info,
            "ncf_info": self.ncf_info,
            "priv_info": self.priv_info,
            "tactile_units": self.network_config.tactile_mlp.units,
            "ncf_units": self.network_config.ncf_mlp.units,
            "ncf_adapt_units": self.network_config.ncf_adapt_mlp.units,
            "tactile_input_shape": (self.tactile_info_dim,),
            "ncf_input_shape": (self.ncf_output_dim,),
            "ncf_proprio_adapt": self.ncf_proprio_adapt,
        }
        self.model = ActorCritic(net_config)
        self.model.to(self.device)
        self.model.eval()

        if self.tactile_info:
            self.obs_shape = (self.obs_shape[0] + self.tactile_info_dim,)

        # ---- PPO Param ----
        self.normalize_input = self.ppo_config["normalize_input"]
        self.normalize_point_cloud = self.ppo_config["normalize_point_cloud"]
        self.obs_mean_std = RunningMeanStd(self.obs_shape).to(self.device)
        self.obs_mean_std.eval()
        self.point_cloud_mean_std = RunningMeanStd(
            3,
        ).to(self.device)
        self.point_cloud_mean_std.eval()

        # ---- DIGIT autoencoder
        # path_checkpoint = self.task_config.ncf.checkpoint_vae
        # self.digit_vae = self.load_digit_autoencoder(self.root_dir, path_checkpoint)
        # self.digit_vae.eval()
        # self.digit_vae.to(self.device)
        # self.reset_digits_fifo()
        # self.reset_eef_fifo()

        # # ---- NDF
        # path_checkpoint = self.task_config.ncf.checkpoint_ndf
        # self.ndf = self.load_ndf(self.root_dir, path_checkpoint)
        # self.ndf.eval()
        # self.ndf.to(self.device)

        # # ---- NCF
        # self.ncf = self.load_ncf(self.root_dir)
        # self.ncf.eval()
        # self.ncf.to(self.device)
        # # self.ncf_viz = NCF_Viz(
        # #     pc_file=self.task_config.ncf.path_pointcloud_object,
        # #     mesh_file=self.task_config.ncf.path_mesh_object,
        # # )

        # # ---- NCF Viz
        # self.viz = NCFViz(off_screen=False, zoom=1.0, window_size=0.70)

        # ---- NCF
        path_checkpoint_vae = self.task_config.ncf.checkpoint_vae
        path_checkpoint_ndf = self.task_config.ncf.checkpoint_ndf
        path_checkpoint_ncf = "{0}/{1}_{2}/model-epoch={3:03d}.ckpt".format(
            self.root_dir,
            self.task_config.ncf.checkpoint_ncf,
            self.task_config.ncf.ncf_arch,
            self.task_config.ncf.ncf_epoch,
        )
        self.digit_vae, self.ndf, self.ncf = self.load_ncf(
            self.task_config.ncf.ncf_arch,
            self.root_dir,
            path_checkpoint_vae,
            path_checkpoint_ndf,
            path_checkpoint_ncf,
        )
        self.digit_vae.eval()
        self.digit_vae.to(self.device)
        self.ncf.eval()
        self.ncf.to(self.device)

        # ---- NCF Viz
        self.viz = NCFViz(off_screen=False, zoom=1.0, window_size=0.70)

        # ---- NCF data
        path_ndf_code_object = self.task_config.ncf.path_ndf_code_object
        self.shape_code = np.load(path_ndf_code_object)
        # self.shape_code = np.zeros((1, 768))
        self.shape_code = torch.tensor(self.shape_code, dtype=torch.float32).unsqueeze(
            0
        )

        path_pointcloud_object = self.task_config.ncf.path_pointcloud_object
        self.ncf_point_cloud = np.load(path_pointcloud_object)
        self.ncf_point_cloud = self.ncf_point_cloud - np.mean(
            self.ncf_point_cloud, axis=0
        )
        self.n_points = len(self.ncf_point_cloud)

        # ---- Contact Filter
        win_sz = 5
        self.contact_buffer = torch.zeros((win_sz, 2000), dtype=torch.float32)

        # ---- loggers
        self.reset_loggers()

        self.stop_episode = False

    def restore(self, fn):
        checkpoint = torch.load(fn)
        self.model.load_state_dict(checkpoint["model"])
        if self.normalize_input:
            self.obs_mean_std.load_state_dict(checkpoint["running_mean_std"])
        if self.normalize_point_cloud:
            self.point_cloud_mean_std.load_state_dict(
                checkpoint["point_cloud_mean_std"]
            )

    def load_digit_autoencoder(self, root_dir, path_checkpoint):
        from algo2.ncf.vae.vae import VAE

        path_checkpoint = os.path.join(root_dir, path_checkpoint)
        vae = VAE.load_from_checkpoint(
            path_checkpoint,
            enc_out_dim=512,
            latent_dim=64,
            input_height=64,
        )
        vae.eval()
        for param in vae.parameters():
            param.requires_grad = False
        return vae

    def load_ndf(self, root_dir, path_checkpoint):
        from algo2.ncf.ndf.ndf import NDF

        path_checkpoint = os.path.join(root_dir, path_checkpoint)
        ndf = NDF(latent_dim=256, return_features=True, sigmoid=True)
        checkpoint = torch.load(path_checkpoint)
        ndf_weights = ndf.state_dict()

        for key in ndf_weights.keys():
            # print(key)
            ndf_weights[key] = checkpoint[key]
        ndf.load_state_dict(ndf_weights)
        for param in ndf.parameters():
            param.requires_grad = False

        return ndf

    def load_ncf(
        self,
        arch,
        root_dir,
        path_checkpoint_vae,
        path_checkpoint_ndf,
        path_checkpoint_ncf,
    ):
        from algo2.ncf.contact_regressor.ncf_mlp import NCF as NCF_mlp
        from algo2.ncf.contact_regressor.ncf_transformer import NCF as NCF_transformer
        from algo2.ncf.config.config import NCF_Params
        from algo2.ncf.pipeline import NCF_Pipeline

        cfg = NCF_Params()
        digit_encoder = self.load_digit_autoencoder(root_dir, path_checkpoint_vae)
        ndf_model = self.load_ndf(root_dir, path_checkpoint_ndf)

        if "mlp" in arch:
            ncf_model = NCF_mlp(cfg)
        elif "transformer" in arch:
            ncf_model = NCF_transformer(cfg)
        else:
            raise NotImplementedError

        pipeline = NCF_Pipeline.load_from_checkpoint(
            path_checkpoint_ncf,
            cfg=cfg,
            digit_vae=digit_encoder,
            ndf=ndf_model,
            ncf=ncf_model,
        )

        for param in pipeline.parameters():
            param.requires_grad = False

        # digit_encoder = self.load_digit_autoencoder(root_dir, path_checkpoint_vae)
        # return digit_encoder, pipeline.ncf
        return pipeline.digit_vae, pipeline.ndf, pipeline.ncf

    def reset_digits_fifo(self):
        self.digit_left_fifo = torch.zeros(
            (1, self.tactile_seq_length, self.tactile_info_embed_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self.digit_right_fifo = torch.zeros(
            (1, self.tactile_seq_length, self.tactile_info_embed_dim),
            dtype=torch.float32,
            device=self.device,
        )

    def reset_eef_fifo(self):
        self.eef_fifo = torch.zeros(
            (1, self.tactile_seq_length, 7),
            dtype=torch.float32,
            device=self.device,
        )
        self.eef_fifo[:, :, -1] = 1.0

    def _digit_encode(self, images):
        images = images.permute(0, 3, 1, 2).unsqueeze(0)
        with torch.no_grad():
            images_hat, digit_embeddings = self.digit_vae(images)

        return digit_embeddings.unsqueeze(0), images_hat

    def _get_digit_embeddings(self, obs_dict):
        digits_imgs_left = obs_dict["digits_left"]
        digits_imgs_right = obs_dict["digits_right"]

        digits_emb_left, left_hat = self._digit_encode(digits_imgs_left)
        digits_emb_right, right_hat = self._digit_encode(digits_imgs_right)

        self.digit_left_fifo = add_to_fifo(self.digit_left_fifo, digits_emb_left)
        self.digit_right_fifo = add_to_fifo(self.digit_right_fifo, digits_emb_right)

        self.img_left = digits_imgs_left[0].cpu().numpy()
        self.img_right = digits_imgs_right[0].cpu().numpy()
        self.img_left_vae = left_hat[0].permute(1, 2, 0).cpu().numpy()
        self.img_right_vae = right_hat[0].permute(1, 2, 0).cpu().numpy()

    def info2buffers(self, obs_dict):
        self.eef_fifo = add_to_fifo(
            self.eef_fifo, obs_dict["obs_ncf"][0:7].unsqueeze(0).unsqueeze(1)
        )
        self._get_digit_embeddings(obs_dict)
        # print(f"mug ncf = {obs_dict['obs_ncf'][0:3]}")

    def _get_ee_deltas(self):
        ee_seq = self.eef_fifo.clone().squeeze(0).cpu().numpy()
        # ee_seq[:, 2] -= 0.02
        # ee_seq[:, 2] -= 0.03
        r0 = Rot.from_quat(ee_seq[0][3:]).as_matrix()
        r1 = Rot.from_quat(ee_seq[1][3:]).as_matrix()
        r2 = Rot.from_quat(ee_seq[2][3:]).as_matrix()
        q_t_1 = Rot.from_matrix(np.matmul(r0, r1.T)).as_quat()
        q_t_2 = Rot.from_matrix(np.matmul(r0, r2.T)).as_quat()
        t_t_1 = ee_seq[0][0:3] - ee_seq[1][0:3]
        t_t_2 = ee_seq[0][0:3] - ee_seq[2][0:3]
        ee_t1 = np.concatenate((t_t_1, q_t_1))
        ee_t2 = np.concatenate((t_t_2, q_t_2))

        ee_seq[:, 0:3] = ee_seq[:, 0:3] / EEF_MAX
        ee_t1[0:3] = ee_t1[0:3] / DELTA_EEF_MAX
        ee_t1[3:] = ee_t1[3:] / DELTA_QUAT_MAX
        ee_t2[0:3] = ee_t2[0:3] / DELTA_EEF_MAX
        ee_t2[3:] = ee_t2[3:] / DELTA_QUAT_MAX

        ee_seq = torch.tensor(ee_seq, dtype=torch.float32).unsqueeze(0).to(self.device)
        ee_t1 = torch.tensor(ee_t1, dtype=torch.float32).unsqueeze(0).to(self.device)
        ee_t2 = torch.tensor(ee_t2, dtype=torch.float32).unsqueeze(0).to(self.device)

        return ee_seq, ee_t1, ee_t2

    def get_pointcloud_data(self):
        idx_points = np.arange(self.n_points)
        n_query_pts = int(self.n_points * 1.0)
        idx_query = np.random.choice(idx_points, size=n_query_pts)
        ncf_query_points = self.ncf_point_cloud[idx_query]
        ncf_query_points = torch.tensor(ncf_query_points, dtype=torch.float32)
        ncf_query_points = ncf_query_points.unsqueeze(0).to(self.device)
        shape_code = self.shape_code.to(self.device)
        return shape_code, ncf_query_points, idx_query

    def _get_pointcloud_data_ndf(self):
        # points for ncf
        path_point_cloud = self.task_config.ncf.path_pointcloud_object
        path_mesh = self.task_config.ncf.path_mesh_object

        ncf_point_cloud = np.load(path_point_cloud)
        ncf_point_cloud = ncf_point_cloud - np.mean(ncf_point_cloud, axis=0)
        n_points = len(ncf_point_cloud)

        idx_points = np.arange(n_points)
        n_query_pts = int(n_points * 1.0)
        idx_query = np.random.choice(idx_points, size=n_query_pts)
        ncf_query_points = ncf_point_cloud[idx_query]

        ncf_query_points = torch.tensor(ncf_query_points, dtype=torch.float32)
        ncf_query_points = ncf_query_points.unsqueeze(0).to(self.device)

        # points for ndf
        object_trimesh = trimesh.load(path_mesh)
        T = np.eye(4)
        T[0:3, 0:3] = Rot.from_euler("xyz", [0, 0, 90], degrees=True).as_matrix()
        object_trimesh = object_trimesh.apply_transform(T)

        ndf_point_cloud = np.array(
            trimesh.sample.sample_surface(object_trimesh, n_points // 2)[0]
        )
        ndf_point_cloud = ndf_point_cloud - np.mean(ndf_point_cloud, axis=0)
        ndf_point_cloud = torch.tensor(ndf_point_cloud, dtype=torch.float32)
        ndf_point_cloud = ndf_point_cloud.unsqueeze(0).to(self.device)

        return ndf_point_cloud, ncf_query_points, idx_query

    def _get_ncf_output(self, obs_dict):
        self.info2buffers(obs_dict)

        digits_emb_left = self.digit_left_fifo.reshape(
            -1, self.tactile_seq_length * self.tactile_info_embed_dim
        )
        digits_emb_right = self.digit_right_fifo.reshape(
            -1, self.tactile_seq_length * self.tactile_info_embed_dim
        )

        ee_seq, ee_t1, ee_t2 = self._get_ee_deltas()

        # (
        #     shape_code,
        #     ncf_query_points,
        #     idx_query,
        # ) = self.get_pointcloud_data()

        ndf_point_cloud, ncf_query_points, idx_query = self._get_pointcloud_data_ndf()
        with torch.no_grad():
            shape_code = self.ndf(ndf_point_cloud)

        inputs_ncf = {
            "digit_emb_left": digits_emb_left,
            "digit_emb_right": digits_emb_right,
            "ee_pose": ee_seq,
            "ee_pose_1": ee_t1,
            "ee_pose_2": ee_t2,
            "shape_emb": shape_code,
            "ncf_query_points": ncf_query_points,
        }
        with torch.no_grad():
            contact = self.ncf(inputs_ncf)

        # d_mug_cupholder = (
        #     torch.norm(self.env.nut_pos - self.env.bolt_pos, dim=1)
        # ) > 0.13
        d_mug_cupholder = self.eef_fifo[0][0][2] > 0.09
        if d_mug_cupholder:
            contact = torch.zeros_like(contact)
        # contact = torch.where(contact > 0.4, 1.0, 0.0)
        contact[contact > 0.4] = 1.0
        contact_gt = contact.clone()

        return contact_gt, contact, idx_query, inputs_ncf

    def log_data(self, obs_dict, inputs_ncf, ncf_contact, idx_query):
        for key in inputs_ncf.keys():
            inputs_ncf[key] = inputs_ncf[key].cpu().numpy()

        joint_pos = self.franka_env.get_joint_positions()
        self.log_joint_pos.append(joint_pos.numpy())
        self.log_obs.append(obs_dict["obs"].cpu().numpy())
        self.log_obs_ncf.append(obs_dict["obs_ncf"].cpu().numpy())
        self.log_obs_ee.append(obs_dict["obs_ee"].cpu().numpy())
        self.log_ncf_inputs.append(inputs_ncf)
        self.log_ncf_contact.append(ncf_contact.cpu().numpy())
        self.log_ncf_idx_query.append(idx_query)
        self.log_digit_left_imgs.append((self.img_left * 255).astype(np.uint8))
        self.log_digit_right_imgs.append((self.img_right * 255).astype(np.uint8))
        self.log_rs_front.append(self.franka_env.rs_front_img)
        self.log_rs_rear.append(self.franka_env.rs_rear_img)

    def reset_loggers(self):
        self.log_joint_pos = []
        self.log_obs = []
        self.log_obs_ncf = []
        self.log_obs_ee = []
        self.log_ncf_inputs = []
        self.log_ncf_contact = []
        self.log_ncf_idx_query = []
        self.log_digit_left_imgs = []
        self.log_digit_right_imgs = []
        self.log_rs_front = []
        self.log_rs_rear = []

    def save_images(self, imgs, path):
        for i, img in enumerate(imgs):
            img = img[:, :, ::-1]
            cv2.imwrite(f"{path}/{i:03d}.png", img)

    def save_logs(self, path_save):
        os.makedirs(path_save, exist_ok=True)
        path_file_save_digit_left = f"{path_save}/real_left/"
        path_file_save_digit_right = f"{path_save}/real_right/"
        path_file_save_rs_front = f"{path_save}/rs_front/"
        path_file_save_rs_rear = f"{path_save}/rs_rear/"
        os.makedirs(path_file_save_digit_left, exist_ok=True)
        os.makedirs(path_file_save_digit_right, exist_ok=True)
        os.makedirs(path_file_save_rs_front, exist_ok=True)
        os.makedirs(path_file_save_rs_rear, exist_ok=True)

        self.save_images(self.log_digit_left_imgs, path_file_save_digit_left)
        self.save_images(self.log_digit_right_imgs, path_file_save_digit_right)
        self.save_images(self.log_rs_front, path_file_save_rs_front)
        self.save_images(self.log_rs_rear, path_file_save_rs_rear)

        log_data = {
            "joint_pos": self.log_joint_pos,
            "obs": self.log_obs,
            "obs_ncf": self.log_obs_ncf,
            "obs_ee": self.log_obs_ee,
            "ncf_inputs": self.log_ncf_inputs,
            "ncf_contact": self.log_ncf_contact,
            "idx_query": self.log_ncf_idx_query,
        }
        np.savez_compressed(f"{path_save}/log_data.npz", **log_data)
        self.reset_loggers()

    def get_policy_inputs(self, obs_dict):
        if self.tactile_info:
            self.info2buffers(obs_dict)
            digits_emb_left = self.digit_left_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )
            digits_emb_right = self.digit_right_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )

            digits_emb = torch.cat([digits_emb_left, digits_emb_right], dim=-1)

            aug_obs = torch.cat([obs_dict["obs"], digits_emb[0]], dim=-1)
            processed_obs = self.obs_mean_std(aug_obs).unsqueeze(0)

        else:
            aug_obs = obs_dict["obs"]
            processed_obs = self.obs_mean_std(obs_dict["obs"]).unsqueeze(0)

        return processed_obs, aug_obs

    def on_press(self, key):
        if key == keyboard.Key.space:
            self.stop_episode = True

    def deploy(self, c):
        # create franka object and connect to robot
        # ----- Robot -------
        self.franka_env = FrankaEnv(
            self.task_config,
            self.action_scale_pos,
            self.action_scale_rot,
            with_realsense=False,
            pointcloud_obj=self.task_config.ncf.path_pointcloud_object,
            use_worldref_pointcloud=self.task_config.rl.world_ref_pointcloud,
        )

        self.viz.init_variables(
            object_mesh_path=self.task_config.ncf.path_mesh_object,
            object_pointcloud=self.task_config.ncf.path_pointcloud_object,
        )

        start_episode = 0
        n_episodes = 15
        f_sample = 30.0

        if self.tactile_info:
            path_save = (
                "/home/fair/carolina/NCF_RL/outputs_policies/paper/proprio_tactile/"
            )
        elif self.ncf_info:
            path_save = "/home/fair/carolina/NCF_RL/outputs_policies/paper/proprio_ncf_{0}_epoch_{1:02d}/".format(
                self.task_config.ncf.ncf_arch,
                self.task_config.ncf.ncf_epoch,
            )
        else:
            path_save = (
                "/home/fair/carolina/NCF_RL/outputs_policies/paper/proprio_only/"
            )
        os.makedirs(path_save, exist_ok=True)

        self.stop_episode = False
        listener = keyboard.Listener(on_press=self.on_press)
        listener.start()

        for ep in range(start_episode, start_episode + n_episodes):
            cprint(f"Starting episode {ep}", "green", attrs=["bold"])

            obs_dict = self.franka_env.reset(ep)
            self.franka_env.read_realsense()

            self.franka_env.print_robot_status()
            self.reset_digits_fifo()
            self.reset_eef_fifo()

            for _ in range(10):
                self.info2buffers(obs_dict)

            step = 0
            path_save_logger = f"{path_save}/test_{ep}/"
            self.stop_episode = False

            while True:
                # for i in range(10):
                self.franka_env.read_realsense()
                # processed_obs, aug_obs, ncf_output = self.get_policy_inputs(obs_dict)
                # input_dict = {
                #     "obs": processed_obs,
                # }
                processed_obs, aug_obs = self.get_policy_inputs(obs_dict)
                # processed_obs = self.obs_mean_std(obs_dict["obs"]).unsqueeze(0)
                point_cloud = None

                if self.priv_info:
                    (
                        ncf_priv_info,
                        pred_contact,
                        idx_query,
                        inputs_ncf,
                    ) = self._get_ncf_output(obs_dict)
                    point_cloud = obs_dict["pointclouds_t"]
                    ncf_mask = torch.where(ncf_priv_info > 0.1, 1.0, 0.0)
                    ncf_mask = ncf_mask.repeat(3, 1, 1).permute(1, 2, 0)
                    point_cloud = point_cloud * ncf_mask
                    point_cloud = point_cloud.unsqueeze(0)
                    obs_dict["pointclouds_t"] = point_cloud

                    # if point_cloud[0].sum() > 0.0:
                    #     print("contact")

                    if self.normalize_point_cloud:
                        point_cloud = self.point_cloud_mean_std(
                            point_cloud.reshape(-1, 3)
                        ).reshape((processed_obs.shape[0], -1, 3))

                input_dict = {
                    "obs": processed_obs,
                    "priv_point_cloud": point_cloud,
                }

                action = self.model.act_inference(input_dict).squeeze()
                action = torch.clamp(action, -1.0, 1.0)
                obs_dict, r, done, info = self.franka_env.step(action)
                # self.franka_env.print_robot_status()

                if not self.ncf_info:
                    (
                        ncf_priv_info,
                        pred_contact,
                        idx_query,
                        inputs_ncf,
                    ) = self._get_ncf_output(obs_dict)

                self.log_data(obs_dict, inputs_ncf, pred_contact, idx_query)

                info["episode"] = ep
                info["step"] = step
                info["policy"] = (
                    "proprioception only"
                    if not self.ncf_info
                    else "proprioception + gt contact"
                )

                if self.stop_episode:
                    print("Stopping episode...")
                    break

                # path_save_step = path_save + "ep_{0:02d}_step_{1:02d}.png".format(
                #     ep, step
                # )
                # path_save_logger_step = path_save + "ep_{0:02d}.npz".format(ep)

                # self.viz.update(
                #     info=info,
                #     rs_img=[self.franka_env.rs_front_img, self.franka_env.rs_rear_img],
                #     digit_left=self.img_left,
                #     digit_right=self.img_right,
                #     vae_left=self.img_left_vae,
                #     vae_right=self.img_right_vae,
                #     contact_vector=pred_contact,
                #     idx_query=idx_query,
                #     image_savepath=None,
                # )

                if done:
                    self.franka_env.print_robot_status()
                    print("")
                    print("is_mug_inserted = {0}".format(info["is_mug_inserted"]))
                    print("is_mug_close = {0}".format(info["is_mug_close"]))
                    print("is_mug_oriented = {0}".format(info["is_mug_oriented"]))
                    print("error pos = {0:04f}".format(info["err_pos"]))
                    print("error quat = {0:04f}".format(info["err_quat"]))
                    print("")
                    break

                step += 1
                time.sleep(1 / f_sample)
            # self.viz.save_imgs()
            self.save_logs(path_save_logger)
        self.viz.close()

NCF_policies/algo2/policy/deploy/deploy.py
carolinahiguera/NCF-Policy
import os
import time
import cv2
import numpy as np
import matplotlib.pyplot as plt
import trimesh
from isaacgymenvs.utils import torch_jit_utils as torch_utils

import torch
from torchcontrol.transform import Rotation as R
from scipy.spatial.transform import Rotation as Rot
from polymetis import RobotInterface
from digit_interface.digit import Digit
import pyrealsense2 as rs

from algo2.utils.misc import tprint


class FrankaEnv:
    def __init__(
        self,
        task_config,
        action_scale_pos,
        action_scale_rot,
        pointcloud_obj,
        with_realsense: bool = False,
        use_worldref_pointcloud: bool = True,
    ):
        print("[robot] configuring robot - setting home pose...")
        self.robot = RobotInterface(ip_address="172.16.0.1", enforce_version=False)
        self.task_config = task_config
        self.action_scale_pos = action_scale_pos
        self.action_scale_rot = action_scale_rot
        # self.home_pose = torch.tensor(self.task_config.randomize.franka_arm_initial_dof_pos)
        self.home_pose = torch.tensor(
            [0.0139, -0.0544, -0.0048, -2.2373, -0.0085, 2.1169, 0.7795]
        )

        # transformations
        # self.ee_offset = torch.tensor([0.5119, 0.0079, 0.19]) # this is the position of the cupholder
        # [0.5319, 0.0079, 0.19]
        self.ee_offset = torch.tensor(
            [0.4931, 0.0041, 0.2222]  # proprio_only
            # [0.50, -0.02, 0.19]
        )  # this is the position of the cupholder
        # [0.4792, 0.0101, 0.2111]
        self.ncf_offset = torch.tensor([0.4931, 0.0041, 0.2122]) + torch.tensor(
            [0.0, -0.02, -0.01]
        )
        # self.ncf_offset = torch.tensor([0.4792, 0.0101, 0.19])
        self.ee_quat_ref = torch.tensor([0.9231, -0.3846, -0.0021, 0.0055])
        ee_quat_ref_sim = torch.tensor([-1.21e-3, 0.99, 2.2e-3, 2.23e-3])
        self.T_real_sim = torch_utils.quat_mul(
            ee_quat_ref_sim, torch_utils.quat_conjugate(self.ee_quat_ref)
        )

        # global variables
        self.digit_sz = 64
        self.h_mug = 0.10
        self.num_envs = 1
        self.device = torch.device("cuda:0")

        self.digit_left = self.config_digit(id="D00011", name="red finger")
        self.digit_right = self.config_digit(id="D00024", name="blue finger")
        self.stabilize_digits()
        self.bg_left, self.bg_right = self.load_bgs()
        print("[robot] configuring DIGIT sensors - ready")

        rs_rear_serial = "152122075373"
        rs_front_serial = "815412070977"
        self.rs_rear_cam = self.config_realsense(rs_rear_serial)
        self.rs_front_cam = self.config_realsense(rs_front_serial)
        self.rs_rear_img = None
        self.rs_front_img = None

        self.pointcloud_obj = np.load(pointcloud_obj)
        self.n_points = self.pointcloud_obj.shape[0]
        self.use_worldref_pointcloud = use_worldref_pointcloud

        self.target_pos = torch.load("target_pos.pt")
        self.target_quat = torch.load("target_quat.pt")

        # self.reset_franka()
        print("[robot] configuring robot - done")

    def load_bgs(self):
        # get path to currrent file
        path = os.path.dirname(os.path.abspath(__file__))
        bg_left = cv2.imread(f"{path}/resources/D00011.png")
        # bg_left = cv2.cvtColor(bg_left, cv2.COLOR_RGB2BGR)
        bg_right = cv2.imread(f"{path}/resources/D00024.png")
        # bg_right = cv2.cvtColor(bg_right, cv2.COLOR_RGB2BGR)
        return bg_left, bg_right

    def config_digit(self, id: str, name: str):
        digit = Digit(id, name)
        digit.connect()
        digit.set_intensity(Digit.LIGHTING_MAX)
        qvga_res = Digit.STREAMS["QVGA"]
        digit.set_resolution(qvga_res)
        fps_30 = Digit.STREAMS["QVGA"]["fps"]["30fps"]
        digit.set_fps(fps_30)
        return digit

    def stabilize_digits(self):
        for _ in range(30):
            self.digit_left.get_frame()
            self.digit_right.get_frame()

    def config_realsense(self, serial):
        # Configure depth and color streams
        pipeline = rs.pipeline()
        config = rs.config()

        config.enable_device(serial)
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

        # Start streaming
        pipeline.start(config)
        return pipeline

    def _subtract_bg(self, img1, img2, offset=0.5):
        img1 = np.int32(img1)
        img2 = np.int32(img2)
        diff = img1 - img2
        diff = diff / 255.0 + offset
        diff = np.clip(diff, 0.0, 1.0)
        diff = cv2.resize(
            diff, (self.digit_sz, self.digit_sz), interpolation=cv2.INTER_AREA
        )
        return diff

    def read_digits(self):
        img_left = self.digit_left.get_frame()
        img_right = self.digit_right.get_frame()
        img_left = self._subtract_bg(img_left, self.bg_left)
        img_right = self._subtract_bg(img_right, self.bg_right)
        img_left = cv2.flip(img_left, 0)
        img_right = cv2.flip(img_right, 0)
        return img_left, img_right

    def read_realsense(self):
        frames = self.rs_front_cam.wait_for_frames()
        color_frame = frames.get_color_frame()
        if color_frame:
            img_realsense = np.asanyarray(color_frame.get_data())
            self.rs_front_img = cv2.cvtColor(img_realsense, cv2.COLOR_RGB2BGR)

        frames = self.rs_rear_cam.wait_for_frames()
        color_frame = frames.get_color_frame()
        if color_frame:
            img_realsense = np.asanyarray(color_frame.get_data())
            self.rs_rear_img = cv2.cvtColor(img_realsense, cv2.COLOR_RGB2BGR)

    def lift_up(self):
        delta_pos = torch.tensor([0.0, 0.0, 0.05])
        self.robot.move_to_ee_pose(position=delta_pos, time_to_go=5.0, delta=True)
        time.sleep(0.1)

    def _start_franka_controller(self):
        Kx = [250, 500, 500, 10, 10, 10]
        Kxd = [25, 5, 1, 7, 7, 7]
        a = self.robot.start_cartesian_impedance(Kx, Kxd)
        time.sleep(0.1)
        print(a)
        print("[robot] cartesian impedance controller initialized")

    def reset_franka(self, ep=-1):
        self.lift_up()
        self.robot.set_home_pose(self.home_pose)
        self.robot.go_home()
        time.sleep(0.5)
        # delta_pos = torch.tensor([0.0, 0.0, -0.05])
        # self.robot.move_to_ee_pose(position=delta_pos, time_to_go=3.0, delta=True)
        # time.sleep(0.5)

        if ep == -1:
            self._randomize_gripper_pose()
        else:
            self.set_init_pose(self.target_pos[ep], self.target_quat[ep])

        time.sleep(0.5)
        self._start_franka_controller()
        self.reset_buffers()
        print("[robot] reset robot")

    def reset_buffers(self):
        digit_seq_len = 1
        digit_sz = 64

        self.obs_ncf_buf = torch.zeros(
            (self.num_envs, 21), device=self.device, dtype=torch.float
        )
        self.obs_buf = torch.zeros(
            (self.num_envs, 21), device=self.device, dtype=torch.float
        )
        self.obs_ee_buf = torch.zeros(
            (self.num_envs, 21), device=self.device, dtype=torch.float
        )
        self.progress_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long
        )
        self.digit_left_buf = torch.zeros(
            (self.num_envs, digit_seq_len, digit_sz, digit_sz, 3),
            device=self.device,
            dtype=torch.float,
        )
        self.digit_right_buf = torch.zeros(
            (self.num_envs, digit_seq_len, digit_sz, digit_sz, 3),
            device=self.device,
            dtype=torch.float,
        )
        self.pointcloud_buf = torch.zeros(
            (self.num_envs, self.n_points, 3),
            device=self.device,
            dtype=torch.float,
        )
        self.obs_dict = {}

    def reset(self, ep=-1):
        self.reset_franka(ep=ep)
        self.get_observations_proprio()
        self.get_observations_tactile()
        self.obs_dict["obs"] = self.obs_buf.to(self.device)
        self.obs_dict["obs_ee"] = self.obs_ee_buf.to(self.device)
        self.obs_dict["obs_ncf"] = self.obs_ncf_buf.to(self.device)
        self.obs_dict["digits_left"] = self.digit_left_buf.to(self.device)
        self.obs_dict["digits_right"] = self.digit_right_buf.to(self.device)
        self.obs_dict["pointclouds_t"] = self.pointcloud_buf.to(self.device)
        return self.obs_dict

    def get_ee_offset(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        offset = torch.tensor([0.4819, 0.07, 0.14])
        ee_pos = ee_pos - offset
        ee_quat = torch_utils.quat_mul(self.T_real_sim, ee_quat)
        cupholder_pos = torch.tensor([0.0, 0.0, 0.0])
        cupholder_quat = torch.tensor([0.0, 0.0, 0.0, 1.0])
        ee_pose = torch.cat([ee_pos, ee_quat, cupholder_pos, cupholder_quat])
        return ee_pose

    def get_ee_pose(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        # ee_pos = ee_pos.to(self.device)
        # ee_quat = ee_quat.to(self.device)
        return ee_pos, ee_quat

    def get_joint_positions(self):
        joint_positions = self.robot.get_joint_positions()
        return joint_positions

    def print_robot_status(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        mug_pose = ee_pos - self.ee_offset
        # mug_pose[2] -= self.h_mug
        mug_quat = torch_utils.quat_mul(self.T_real_sim, ee_quat)
        info_string = (
            f"[robot] ee_pos = {ee_pos} \n"
            f"[robot] ee_quat = {ee_quat} \n"
            f"[robot] mug_pos (policy) = {mug_pose} \n"
            f"[robot] mug_quat (policy) = {mug_quat} \n"
            f""
        )
        tprint(info_string)

    def _xyzquat_to_tf_numpy(self, position_quat: np.ndarray) -> np.ndarray:
        """
        convert [x, y, z, qx, qy, qz, qw] to 4 x 4 transformation matrices
        """
        # try:
        position_quat = np.atleast_2d(position_quat)  # (N, 7)
        N = position_quat.shape[0]
        T = np.zeros((N, 4, 4))
        T[:, 0:3, 0:3] = Rot.from_quat(position_quat[:, 3:]).as_matrix()
        T[:, :3, 3] = position_quat[:, :3]
        T[:, 3, 3] = 1
        # except ValueError:
        #     print("Zero quat error!")
        return T.squeeze()

    def get_observations_proprio(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        cupholder_pos = torch.tensor([0.0, 0.0, 0.0])
        cupholder_quat = torch.tensor([0.0, 0.0, 0.0, 1.0])
        ee_quat = torch_utils.quat_mul(self.T_real_sim, ee_quat)
        # ee_quat_ncf = torch_utils.quat_mul(torch.tensor([0, 0, 1, 0]), ee_quat)
        obs_tensors = [ee_pos - self.ee_offset, ee_quat]
        obs_ncf_tensors = [ee_pos - self.ncf_offset, ee_quat]
        obs_ee_tensors = [ee_pos, ee_quat]
        # print("ee_ncf = ", ee_pos - self.ncf_offset)

        # transform pointcloud
        if self.use_worldref_pointcloud:
            mug_pos = ee_pos - self.ee_offset
            object_pose = torch.cat((mug_pos, ee_quat), dim=-1).unsqueeze(0)
            object_pose = self._xyzquat_to_tf_numpy(object_pose.cpu().numpy())
            object_pc_i = trimesh.points.PointCloud(self.pointcloud_obj.copy())
            object_pc_i.apply_transform(object_pose)
            coords = object_pc_i.vertices
        else:
            coords = self.pointcloud_obj.copy()

        if self.obs_buf.sum() == 0:
            curr_obs = torch.cat(obs_tensors, dim=-1)
            self.obs_buf = torch.cat((curr_obs, curr_obs, curr_obs), dim=-1)
            curr_obs = torch.cat(obs_ncf_tensors, dim=-1)
            self.obs_ncf_buf = torch.cat((curr_obs, curr_obs, curr_obs), dim=-1)
            curr_obs = torch.cat(obs_ee_tensors, dim=-1)
            self.obs_ee_buf = torch.cat((curr_obs, curr_obs, curr_obs), dim=-1)
        else:
            curr_obs = torch.cat(obs_tensors, dim=-1)
            prev_obs_buf = self.obs_buf[0:14].clone()
            self.obs_buf = torch.cat((curr_obs, prev_obs_buf), dim=-1)
            curr_obs = torch.cat(obs_ncf_tensors, dim=-1)
            prev_obs_buf = self.obs_ncf_buf[0:14].clone()
            self.obs_ncf_buf = torch.cat((curr_obs, prev_obs_buf), dim=-1)
            curr_obs = torch.cat(obs_ee_tensors, dim=-1)
            prev_obs_buf = self.obs_ee_buf[0:14].clone()
            self.obs_ee_buf = torch.cat((curr_obs, prev_obs_buf), dim=-1)

        self.pointcloud_buf = torch.tensor(
            coords, device=self.device, dtype=torch.float
        )

        # print(f"[robot] obs_buf = {self.obs_buf[0:3]}")
        # print(f"[robot] obs_ncf_buf = {self.obs_ncf_buf[0:3]}")

    def get_observations_tactile(self):
        digit_seq_len = 1
        digit_sz = 64
        img_left, img_right = self.read_digits()

        self.digit_left_buf = torch.tensor(img_left, dtype=torch.float).unsqueeze(0)
        self.digit_right_buf = torch.tensor(img_right, dtype=torch.float).unsqueeze(0)

    def _randomize_gripper_pose(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        ctrl_target_fingertip_midpoint_pos = ee_pos
        fingertip_midpoint_pos_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32) - 0.5
        )  # [-1, 1]
        fingertip_midpoint_pos_noise = fingertip_midpoint_pos_noise @ torch.diag(
            torch.tensor(self.task_config.randomize.fingertip_midpoint_pos_noise)
        )

        # fingertip_midpoint_pos_noise[0, 2] = 0.0
        ctrl_target_fingertip_midpoint_pos += fingertip_midpoint_pos_noise.squeeze()

        x, y, z = torch_utils.get_euler_xyz(ee_quat.unsqueeze(0))
        ctrl_target_fingertip_midpoint_euler = torch.cat([x, y, z])
        fingertip_midpoint_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32) - 0.5
        )  # [-1, 1]
        fingertip_midpoint_rot_noise = fingertip_midpoint_rot_noise @ torch.diag(
            torch.tensor(self.task_config.randomize.fingertip_midpoint_rot_noise)
        )
        ctrl_target_fingertip_midpoint_euler += fingertip_midpoint_rot_noise.squeeze()
        ctrl_target_fingertip_midpoint_quat = (
            torch_utils.quat_from_euler_xyz(
                ctrl_target_fingertip_midpoint_euler[0],
                ctrl_target_fingertip_midpoint_euler[1],
                ctrl_target_fingertip_midpoint_euler[2],
            )
            * -1.0
        )

        self.robot.move_to_ee_pose(
            position=ctrl_target_fingertip_midpoint_pos,
            orientation=ctrl_target_fingertip_midpoint_quat,
            time_to_go=3.0,
            delta=False,
        )
        # return ctrl_target_fingertip_midpoint_pos, ctrl_target_fingertip_midpoint_quat

    def set_init_pose(
        self, ctrl_target_fingertip_midpoint_pos, ctrl_target_fingertip_midpoint_quat
    ):
        self.robot.move_to_ee_pose(
            position=ctrl_target_fingertip_midpoint_pos,
            orientation=ctrl_target_fingertip_midpoint_quat,
            time_to_go=3.0,
            delta=False,
        )

    def _apply_actions_as_ctrl_targets(self, actions, do_scale=False):
        actions = actions.cpu()
        """Apply actions from policy as position/rotation targets."""
        ee_pos, ee_quat = self.get_ee_pose()
        ee_quat = torch_utils.quat_mul(self.T_real_sim, ee_quat)
        # Interpret actions as tar get pos displacements and set pos target
        pos_actions = actions[0:3] * self.action_scale_pos
        if do_scale:
            pos_actions = pos_actions @ torch.diag(
                torch.tensor(self.task_config.rl.pos_action_scale)
            )

        # Interpret actions as target rot (axis-angle) displacements
        rot_actions = actions[3:6] * self.action_scale_rot
        if do_scale:
            rot_actions = rot_actions @ torch.diag(
                torch.tensor(self.task_config.rl.rot_action_scale)
            )
        # Convert to quat and set rot target
        angle = torch.norm(rot_actions, p=2, dim=-1)
        axis = rot_actions / angle.unsqueeze(-1)
        rot_actions_quat = torch_utils.quat_from_angle_axis(angle, axis)
        if self.task_config.rl.clamp_rot:
            rot_actions_quat = torch.where(
                angle.unsqueeze(-1).repeat(1, 4) > self.task_config.rl.clamp_rot_thresh,
                rot_actions_quat,
                torch.tensor([0.0, 0.0, 0.0, 1.0]),
            )
        rot_actions_quat = rot_actions_quat.squeeze()
        rot_actions_quat = torch_utils.quat_mul(rot_actions_quat, ee_quat)

        # print(f"pos_actions = {pos_actions}")
        # print(f"rots_actions = {rot_actions}")
        # print("")

        # pos_actions[torch.abs(pos_actions) < 0.005] *= 5.0
        # pos_actions[0] = 0.0
        # pos_actions[1] = 0.0
        self.ctrl_target_ee_pos = ee_pos + pos_actions

        self.ctrl_target_ee_quat = torch_utils.quat_mul(
            torch_utils.quat_conjugate(self.T_real_sim), rot_actions_quat
        )
        self.ctrl_target_ee_quat = torch_utils.normalize(self.ctrl_target_ee_quat)

        # self.ctrl_target_ee_quat = torch.tensor([0.9231, -0.3846, -0.0021, 0.0055])

        # send actions to robot
        try:
            res = self.robot.update_desired_ee_pose(
                position=self.ctrl_target_ee_pos, orientation=self.ctrl_target_ee_quat
            )
        except:
            self._start_franka_controller()
        # print(res)
        time.sleep(0.1)
        # except:
        #     self._start_franka_controller()

    def check_mug_far(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        if (ee_pos[0] < 0.40 and ee_pos[2] < 0.25) or (
            torch.abs(ee_pos[1]) > 0.05 and ee_pos[0] > 0.50 and ee_pos[2] > 0.30
        ):
            return True
        else:
            return False

    def check_mug_in_cupholder(self):
        ee_pos, ee_quat = self.robot.get_ee_pose()
        mug_pos = ee_pos - self.ee_offset
        cupholder_pos = torch.tensor([0.0, 0.01, 0.02])
        err_pos = torch.norm(cupholder_pos - mug_pos, p=2, dim=-1)
        err_quat = torch.norm(self.ee_quat_ref - ee_quat, p=2, dim=-1)
        return err_pos, err_quat

    def step(self, actions):
        self.actions = actions.clone().to(self.device)
        self._apply_actions_as_ctrl_targets(actions=self.actions, do_scale=True)

        # get observations
        self.get_observations_proprio()
        self.get_observations_tactile()
        self.obs_dict["obs"] = self.obs_buf.to(self.device)
        self.obs_dict["obs_ncf"] = self.obs_ncf_buf.to(self.device)
        self.obs_dict["obs_ee"] = self.obs_ee_buf.to(self.device)
        self.obs_dict["digits_left"] = self.digit_left_buf.to(self.device)
        self.obs_dict["digits_right"] = self.digit_right_buf.to(self.device)
        self.obs_dict["pointcloud_t"] = self.pointcloud_buf.to(self.device)

        err_pos, err_quat = self.check_mug_in_cupholder()
        r = (
            self.task_config.rl.keypoint_reward_scale * err_pos
            + self.task_config.rl.orientation_penalty_scale * err_quat
        )

        is_mug_close = err_pos < 0.021  # self.task_config.rl.close_error_thresh
        is_mug_oriented = err_quat < 0.1
        is_mug_inserted = is_mug_close * is_mug_oriented
        # is_mug_far = self.check_mug_far()
        is_mug_far = False

        # is_last_step = self.progress_buf[0] == self.task_config.rl.max_episode_length - 1
        is_last_step = self.progress_buf[0] == 100 - 1
        # done = (
        #     is_last_step.item() or is_mug_close.item() or is_mug_far
        # )  # TODO: should be is_mug_inserted
        done = is_last_step.item() or is_mug_far

        # done = False

        info = {
            "is_mug_close": is_mug_close,
            "is_mug_oriented": is_mug_oriented,
            "is_mug_inserted": is_mug_inserted,
            "err_pos": err_pos,
            "err_quat": err_quat,
        }

        self.progress_buf[0] += 1

        if done:
            self.robot.terminate_current_policy()

        return self.obs_dict, r, done, info

NCF_policies/algo2/policy/deploy/franka.py
carolinahiguera/NCF-Policy
import numpy as np
import pyvista as pv
from matplotlib import cm
from pyvistaqt import BackgroundPlotter
import trimesh
import tkinter as tk
import queue
import copy
import cv2
from PIL import Image
from scipy.spatial.transform import Rotation as Rot
from vedo import trimesh2vedo
from vedo import Plotter, Mesh, Points, Text2D


class NCFViz:
    def __init__(
        self, off_screen: bool = False, zoom: float = 1.0, window_size: int = 0.5
    ):
        pv.global_theme.multi_rendering_splitting_position = 0.7
        """
            subplot(0, 0) main viz
            subplot(0, 1): tactile image viz
            subplot(1, 1): tactile codebook viz 
        """
        shape, row_weights, col_weights = ((2, 3), [0.8, 0.2], [0.3, 0.3, 0.4])
        groups = [(0, 0), (0, 1), (0, 2), (1, 0), (1, np.s_[1:2])]

        w, h = tk.Tk().winfo_screenwidth(), tk.Tk().winfo_screenheight()

        self.plotter = BackgroundPlotter(
            title="NCF",
            lighting="three lights",
            window_size=(int(w * window_size), int(h * window_size)),
            off_screen=off_screen,
            shape=shape,
            row_weights=row_weights,
            col_weights=col_weights,
            groups=groups,
            border_color="white",
            toolbar=False,
            menu_bar=False,
            auto_update=True,
        )
        self.zoom = zoom

        self.viz_queue = queue.Queue(1)
        self.plotter.add_callback(self.update_viz, interval=100)  # 50
        self.pause = False
        self.font_size = int(20 * window_size)

        self.cam_obj_1 = dict(
            position=(0.192824, 4.66691e-3, -0.249805),
            focal_point=(-4.82266e-3, -2.87693e-3, 5.10574e-7),
            viewup=(0.655673, -0.564232, 0.501733),
            distance=0.318628,
            clipping_range=(0.143946, 0.500430),
        )
        self.cam_obj_2 = dict(
            position=(-0.100776, -0.139449, -0.271581),
            focal_point=(3.92765e-3, -2.10310e-3, -3.81815e-3),
            viewup=(-0.904989, 0.398268, 0.149592),
            distance=0.318628,
            clipping_range=(0.165546, 0.532542),
        )

    def init_variables(self, object_mesh_path, object_pointcloud):
        self.object_pointcloud = np.load(object_pointcloud)
        self.n_points = self.object_pointcloud.shape[0]

        object_trimesh = trimesh.load_mesh(object_mesh_path)
        T = np.eye(4)
        T[0:3, 0:3] = Rot.from_euler("xyz", [0, 0, 90], degrees=True).as_matrix()
        object_trimesh = object_trimesh.apply_transform(T)

        # self.object_mesh = trimesh2vedo(object_trimesh).clone()
        # self.object_mesh.subdivide(n=5, method=2)
        self.object_mesh = pv.wrap(object_trimesh)
        self.object_cloud = pv.PolyData(self.object_pointcloud)

        self.plotter.subplot(0, 0)
        self.plotter.add_text(
            "Realsense camera",
            position="upper_left",
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="rs_title",
        )
        self.plotter.camera_position = "xy"
        self.plotter.camera.zoom(2)
        # self.plotter.camera_set = True

        self.plotter.subplot(0, 1)
        self.plotter.camera.Zoom(1)
        self.plotter.add_text(
            "DIGIT and VAE images",
            position="upper_left",
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="digit_title",
        )

        self.plotter.subplot(0, 2)
        self.plotter.camera.zoom(1.0)
        self.plotter.add_text(
            "NCF contact patch estimation",
            position="upper_left",
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="ncf_title",
        )
        # self.plotter.camera_position = "yx"
        # self.plotter.camera.azimuth = 30
        # self.plotter.camera.Zoom(3.0)
        self.plotter.camera.position = self.cam_obj_1["position"]
        self.plotter.camera.focal_point = self.cam_obj_1["focal_point"]
        self.plotter.camera.up = self.cam_obj_1["viewup"]
        self.plotter.camera.distance = self.cam_obj_1["distance"]
        self.plotter.camera.clipping_range = self.cam_obj_1["clipping_range"]
        self.plotter.camera.Zoom(1.0)
        self.plotter.camera_set = True

        self.images = {"im": [], "path": []}
        self.image_plane = None

    def update_viz(self):
        if self.viz_queue.qsize():
            (
                info,
                rs_img,
                digit_left,
                digit_right,
                vae_left,
                vae_right,
                contact_vector,
                idx_query,
                image_savepath,
            ) = self.viz_queue.get()

            self.viz_rs_image(rs_img)
            self.viz_digit_images(digit_left, digit_right, vae_left, vae_right)
            self.viz_ncf_contact(contact_vector, idx_query)
            # self.viz_info(info)

            if image_savepath:
                self.images["im"].append(self.plotter.screenshot())
                self.images["path"].append(image_savepath)

            self.viz_queue.task_done()

    def update(
        self,
        info,
        rs_img,
        digit_left,
        digit_right,
        vae_left,
        vae_right,
        contact_vector,
        idx_query,
        image_savepath,
    ):
        if self.viz_queue.full():
            self.viz_queue.get()
        self.viz_queue.put(
            (
                info,
                rs_img,
                digit_left,
                digit_right,
                vae_left,
                vae_right,
                contact_vector,
                idx_query,
                image_savepath,
            ),
            block=False,
        )

    def viz_rs_image(self, rs_imgs):
        front = rs_imgs[0]
        back = rs_imgs[1]

        white = [255, 255, 255]
        front = cv2.copyMakeBorder(
            front, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )
        back = cv2.copyMakeBorder(
            back, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )
        rs_img = np.concatenate((front, back), axis=0)

        s = 10.0
        # if self.image_plane is None:
        image_plane = pv.Plane(
            direction=(0.0, 0.0, 1.0),
            i_size=rs_img.shape[1] * s,
            j_size=rs_img.shape[0] * s,
            i_resolution=rs_img.shape[1] - 1,
            j_resolution=rs_img.shape[0] - 1,
        )
        # self.image_plane
        image_plane.points[:, -1] = 0.25

        self.plotter.subplot(0, 0)
        # self.plotter.view_xy()
        # self.plotter.camera.Zoom(2.0)

        # image_tex = pv.numpy_to_texture(cv2.flip(rs_img, 1))
        image_tex = pv.numpy_to_texture(rs_img)
        self.plotter.add_mesh(
            image_plane,
            texture=image_tex,
            smooth_shading=False,
            show_scalar_bar=False,
            name="rs_img",
            render=False,
        )
        self.plotter.view_xy()
        self.plotter.camera.Zoom(1.5)

    def viz_digit_images(self, digit_left, digit_right, vae_left, vae_right):
        img_left = (cv2.resize(digit_left, (240, 320)) * 255).astype(np.uint8)
        img_right = (cv2.resize(digit_right, (240, 320)) * 255).astype(np.uint8)
        img_left_vae = (cv2.resize(vae_left, (240, 320)) * 255).astype(np.uint8)
        img_right_vae = (cv2.resize(vae_right, (240, 320)) * 255).astype(np.uint8)

        white = [255, 255, 255]
        img_left = cv2.copyMakeBorder(
            img_left, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )
        img_right = cv2.copyMakeBorder(
            img_right, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )
        img_left_vae = cv2.copyMakeBorder(
            img_left_vae, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )
        img_right_vae = cv2.copyMakeBorder(
            img_right_vae, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=white
        )

        img_left = cv2.putText(
            img_left,
            "Left",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
        img_right = cv2.putText(
            img_right,
            "Right",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
        img_left_vae = cv2.putText(
            img_left_vae,
            "Left VAE",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
        img_right_vae = cv2.putText(
            img_right_vae,
            "Right VAE",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )

        img_org = np.concatenate((img_left, img_right), axis=1)
        img_vae = np.concatenate((img_left_vae, img_right_vae), axis=1)
        img_digits = np.concatenate((img_org, img_vae), axis=0)
        img_digits = cv2.flip(img_digits, 1)

        s = 1.0
        image_plane = pv.Plane(
            direction=(0.0, 0.0, 1.0),
            i_size=img_digits.shape[1] * s,
            j_size=img_digits.shape[0] * s,
            i_resolution=img_digits.shape[1] - 1,
            j_resolution=img_digits.shape[0] - 1,
        )
        # self.image_plane
        image_plane.points[:, -1] = 0.25

        self.plotter.subplot(0, 1)
        image_tex = pv.numpy_to_texture(img_digits)
        self.plotter.add_mesh(
            image_plane,
            texture=image_tex,
            smooth_shading=False,
            show_scalar_bar=False,
            name="digit_img",
            render=False,
        )
        self.plotter.view_xy()
        self.plotter.camera.zoom(1.5)

    # def viz_ncf_contact(self, contact_vector, idx_query):
    #     pred_full = np.zeros((self.n_points))
    #     pred_full[idx_query] = contact_vector.cpu().numpy()

    #     pc_pred = Points(self.object_pointcloud, r=5)
    #     pc_pred = pc_pred.cmap(
    #         "plasma",
    #         pred_full,
    #         vmin=0.0,
    #         vmax=1.0,
    #     )

    #     mesh_pred = self.object_mesh.clone()
    #     mesh_pred = mesh_pred.interpolate_data_from(
    #         pc_pred, n=3, on="points", kernel="shepard"
    #     ).cmap("plasma", vmin=0.0, vmax=1.0)

    #     plt = Plotter(
    #         shape=[1, 2],
    #         axes=0,
    #         sharecam=False,
    #         title="NCF mugs",
    #         offscreen=True,
    #     )

    #     pred_name = Text2D("NCF - view 1", s=0.9)
    #     plt.show([mesh_pred, pc_pred, pred_name], at=0, camera=self.cam_obj_1)
    #     pred_name = Text2D("NCF - view 2", s=0.9)
    #     plt.show([mesh_pred, pc_pred, pred_name], at=1, camera=self.cam_obj_2)
    #     img_ncf = plt.screenshot(asarray=True)
    #     plt.close()
    #     img_ncf = cv2.flip(img_ncf, 1)
    #     # increase resolution

    #     scale = 1.3
    #     new_shape = (int(img_ncf.shape[1] * scale), int(img_ncf.shape[0] * scale))
    #     img_ncf = cv2.resize(img_ncf, new_shape, interpolation=cv2.INTER_AREA)

    #     s = 1.0
    #     image_plane = pv.Plane(
    #         direction=(0.0, 0.0, 1.0),
    #         i_size=img_ncf.shape[1] * s,
    #         j_size=img_ncf.shape[0] * s,
    #         i_resolution=img_ncf.shape[1] - 1,
    #         j_resolution=img_ncf.shape[0] - 1,
    #     )
    #     # self.image_plane
    #     image_plane.points[:, -1] = 0.25

    #     self.plotter.subplot(0, 2)
    #     image_tex = pv.numpy_to_texture(img_ncf)
    #     self.plotter.add_mesh(
    #         image_plane,
    #         texture=image_tex,
    #         smooth_shading=False,
    #         show_scalar_bar=False,
    #         name="ncf_img",
    #         render=False,
    #     )
    #     self.plotter.view_xy()
    #     # self.plotter.camera.zoom(1.6)

    def viz_ncf_contact(self, contact_vector, idx_query):
        pred_full = np.zeros((self.n_points))
        pred_full[idx_query] = contact_vector.cpu().numpy()

        contact_cloud = copy.deepcopy(self.object_cloud)
        contact_cloud["val"] = pred_full

        ncf_contact_mesh = self.object_mesh.interpolate(
            contact_cloud,
            # n_points=10,
            # sharpness=1.0,
            sharpness=2.0,
            strategy="null_value",
            radius=self.object_mesh.length / 15,
        )
        dargs = dict(
            cmap=cm.get_cmap("plasma"),
            scalars="val",
            interpolate_before_map=True,
            ambient=0.3,
            opacity=1.0,
            clim=[0.0, 1.0],
            show_scalar_bar=False,
            silhouette=False,
            reset_camera=False,
        )
        self.plotter.subplot(0, 2)
        self.plotter.add_mesh(ncf_contact_mesh, **dargs)
        self.plotter.add_mesh(contact_cloud, **dargs)

    def viz_info(self, info):
        self.plotter.subplot(1, 0)
        self.plotter.add_text(
            "Test {0}".format(info["episode"]),
            position=(50, 90),
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="episode text",
            render=True,
        )
        self.plotter.add_text(
            "Step {0}".format(info["step"]),
            position=(50, 40),
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="step text",
            render=True,
        )
        self.plotter.subplot(1, 1)
        self.plotter.add_text(
            "Policy: \n {0}".format(info["policy"]),
            position=(50, 70),
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="policy text",
            render=True,
        )

        self.plotter.subplot(1, 2)
        self.plotter.add_text(
            "Mug close to cupholder: ",
            position=(50, 90),
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="mug_close text",
            render=True,
        )
        self.plotter.add_text(
            "Yes" if info["is_mug_close"] else "No",
            # "{0}".format(info["is_mug_close"]),
            position=(370, 90),
            color="green" if info["is_mug_close"] else "red",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="mug_close result",
            render=True,
        )

        self.plotter.subplot(1, 2)
        self.plotter.add_text(
            "Mug correctly oriented: ",
            position=(50, 40),
            color="black",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="mug_orn text",
            render=True,
        )
        self.plotter.add_text(
            "Yes" if info["is_mug_oriented"] else "No",
            # "{0}".format(info["is_mug_oriented"]),
            position=(370, 40),
            color="green" if info["is_mug_oriented"] else "red",
            shadow=True,
            font="times",
            font_size=self.font_size,
            name="mug_orn result",
            render=True,
        )

    def save_imgs(self):
        if len(self.images):
            print("Saving images...")
            for im, path in zip(self.images["im"], self.images["path"]):
                im = Image.fromarray(im.astype("uint8"), "RGB")
                im.save(path)
            self.images = {"im": [], "path": []}

    def close(self):
        self.plotter.close()

NCF_policies/algo2/policy/deploy/ncf_viz.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# In-Hand Object Rotation via Rapid Motor Adaptation
# https://arxiv.org/abs/2210.04887
# Copyright (c) 2022 Haozhi Qi
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class MLP(nn.Module):
    def __init__(self, units, input_size, with_last_activation=True):
        super(MLP, self).__init__()
        # use with_last_activation=False when we need the network to output raw values before activation
        layers = []
        for output_size in units:
            layers.append(nn.Linear(input_size, output_size))
            layers.append(nn.ELU())
            input_size = output_size
        if not with_last_activation:
            layers.pop()
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return self.mlp(x)


# class ProprioAdaptTConv(nn.Module):
#     def __init__(self):
#         super(ProprioAdaptTConv, self).__init__()
#         self.channel_transform = nn.Sequential(
#             nn.Linear(16 + 16, 32),
#             nn.ReLU(inplace=True),
#             nn.Linear(32, 32),
#             nn.ReLU(inplace=True),
#         )
#         self.temporal_aggregation = nn.Sequential(
#             nn.Conv1d(32, 32, (9,), stride=(2,)),
#             nn.ReLU(inplace=True),
#             nn.Conv1d(32, 32, (5,), stride=(1,)),
#             nn.ReLU(inplace=True),
#             nn.Conv1d(32, 32, (5,), stride=(1,)),
#             nn.ReLU(inplace=True),
#         )
#         self.low_dim_proj = nn.Linear(32 * 3, 8)

#     def forward(self, x):
#         x = self.channel_transform(x)  # (N, 50, 32)
#         x = x.permute((0, 2, 1))  # (N, 32, 50)
#         x = self.temporal_aggregation(x)  # (N, 32, 3)
#         x = self.low_dim_proj(x.flatten(1))
#         return x


class ActorCritic(nn.Module):
    def __init__(self, kwargs):
        nn.Module.__init__(self)
        actions_num = kwargs.pop("actions_num")
        self.proprio_input_shape = kwargs.pop("proprio_input_shape")[0]
        self.units = kwargs.pop("actor_units")
        self.tactile_info = kwargs.pop("tactile_info")
        self.ncf_info = kwargs.pop("ncf_info")
        self.tactile_units = kwargs.pop("tactile_units")
        self.ncf_units = kwargs.pop("ncf_units")
        self.ncf_adapt_units = kwargs.pop("ncf_adapt_units")
        tactile_input_shape = kwargs.pop("tactile_input_shape")
        ncf_input_shape = kwargs.pop("ncf_input_shape")
        self.ncf_proprio_adapt = kwargs.pop("ncf_proprio_adapt")

        # mlp_input_shape = input_shape[0]
        out_size = self.units[-1]

        # self.priv_info = kwargs['priv_info']
        # self.priv_info_stage2 = kwargs['proprio_adapt']

        if self.tactile_info:
            self.tactile_mlp = MLP(
                units=self.tactile_units, input_size=tactile_input_shape[0]
            )
            mlp_input_shape = self.proprio_input_shape + self.tactile_units[-1]
        elif self.ncf_info:
            # self.ncf_mlp = MLP(units=self.ncf_units, input_size=ncf_input_shape[0])
            self.ncf_mlp = MLP(units=self.ncf_units, input_size=3)

            if self.ncf_proprio_adapt:
                # self.adapt_ncf = MLP(
                #     units=self.ncf_units, input_size=ncf_input_shape[0]
                # )
                self.adapt_ncf = MLP(units=self.ncf_adapt_units, input_size=3)

            mlp_input_shape = self.proprio_input_shape + self.ncf_units[-1]
        else:
            mlp_input_shape = self.proprio_input_shape

        self.actor_mlp = MLP(units=self.units, input_size=mlp_input_shape)
        self.value = torch.nn.Linear(out_size, 1)
        self.mu = torch.nn.Linear(out_size, actions_num)
        self.sigma = nn.Parameter(
            torch.zeros(actions_num, requires_grad=True, dtype=torch.float32),
            requires_grad=True,
        )

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                fan_out = m.kernel_size[0] * m.out_channels
                m.weight.data.normal_(mean=0.0, std=np.sqrt(2.0 / fan_out))
                if getattr(m, "bias", None) is not None:
                    torch.nn.init.zeros_(m.bias)
            if isinstance(m, nn.Linear):
                if getattr(m, "bias", None) is not None:
                    torch.nn.init.zeros_(m.bias)
        nn.init.constant_(self.sigma, 0)

    @torch.no_grad()
    def act(self, obs_dict):
        # used specifically to collection samples during training
        # it contains exploration so needs to sample from distribution
        mu, logstd, value, _, _ = self._actor_critic(obs_dict)
        sigma = torch.exp(logstd)
        distr = torch.distributions.Normal(mu, sigma)
        selected_action = distr.sample()
        result = {
            "neglogpacs": -distr.log_prob(selected_action).sum(
                1
            ),  # self.neglogp(selected_action, mu, sigma, logstd),
            "values": value,
            "actions": selected_action,
            "mus": mu,
            "sigmas": sigma,
        }
        return result

    def ncf_pointcloud_embedding(self, point_cloud_info: torch.Tensor) -> torch.Tensor:
        pcs = self.ncf_mlp(point_cloud_info)
        pcs_filtered = torch.ones_like(pcs) * -1000.0

        a = torch.nonzero(point_cloud_info)
        if len(a) == 0:
            # return torch.zeros((pcs.shape[0], pcs.shape[2])).float().to(pcs.device)
            return torch.zeros((pcs.shape[0], pcs.shape[2])).float().to(pcs.device)
        b = a[:, [0, 1]].cpu().numpy()
        pcs_filtered[b[:, 0], b[:, 1]] = self.ncf_mlp(point_cloud_info)[
            b[:, 0], b[:, 1]
        ]
        emb = torch.max(pcs_filtered, 1)[0]
        emb[emb == -1000.0] = 0.0
        return emb

        # pc_idx = a[:, 0].unique()
        # if len(pc_idx) > 2:
        #     print("aqui")
        # for i in pc_idx:
        #     pc_contact_points = a[torch.where(a[:, 0] == i)[0], 1].unique()
        #     emb[i] = torch.max(pcs[i, pc_contact_points], 0)[0]

        # return torch.max(pcs, 1)[0]
        # return emb

    def ncf_adapt_pointcloud_embedding(
        self, point_cloud_info: torch.Tensor
    ) -> torch.Tensor:
        pcs = self.adapt_ncf(point_cloud_info)
        return torch.max(pcs, 1)[0]

    @torch.no_grad()
    def act_inference(self, obs_dict):
        # used for testing
        mu, logstd, value, _, _ = self._actor_critic(obs_dict)
        return mu

    def _actor_critic(self, obs_dict):
        # obs = obs_dict["obs"]
        # tactile_input = obs_dict["tactile_inputs"]
        # ncf_gt_input = obs_dict["ncf_gt_inputs"]
        # ncf_pred_input = obs_dict["ncf_pred_inputs"]
        obs = obs_dict["obs"][:, 0 : self.proprio_input_shape]
        ncf_emb, ncf_emb_gt = None, None

        if self.tactile_info:
            tactile_input = obs_dict["obs"][:, self.proprio_input_shape :]
            tactile = self.tactile_mlp(tactile_input)
            obs = torch.cat([obs, tactile], dim=-1)

        elif self.ncf_info:
            priv_point_cloud = obs_dict["priv_point_cloud"]
            ncf_emb_gt = self.ncf_pointcloud_embedding(priv_point_cloud)

            if self.ncf_proprio_adapt:
                ncf_point_cloud = obs_dict["ncf_point_cloud"]
                ncf_emb = self.ncf_adapt_pointcloud_embedding(ncf_point_cloud)
                obs = torch.cat([obs, ncf_emb], dim=-1)
            else:
                obs = torch.cat([obs, ncf_emb_gt], dim=-1)

        x = self.actor_mlp(obs)
        value = self.value(x)
        mu = self.mu(x)
        sigma = self.sigma
        return mu, mu * 0 + sigma, value, ncf_emb, ncf_emb_gt

    def forward(self, input_dict):
        prev_actions = input_dict.get("prev_actions", None)
        rst = self._actor_critic(input_dict)
        mu, logstd, value, extrin, extrin_gt = rst
        sigma = torch.exp(logstd)
        distr = torch.distributions.Normal(mu, sigma)
        entropy = distr.entropy().sum(dim=-1)
        prev_neglogp = -distr.log_prob(prev_actions).sum(1)
        result = {
            "prev_neglogp": torch.squeeze(prev_neglogp),
            "values": value,
            "entropy": entropy,
            "mus": mu,
            "sigmas": sigma,
            "extrin": extrin,
            "extrin_gt": extrin_gt,
        }
        return result

NCF_policies/algo2/policy/models/models.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# In-Hand Object Rotation via Rapid Motor Adaptation
# https://arxiv.org/abs/2210.04887
# Copyright (c) 2022 Haozhi Qi
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
# Based on: IsaacGymEnvs
# Copyright (c) 2018-2022, NVIDIA Corporation
# Licence under BSD 3-Clause License
# https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/
# --------------------------------------------------------

import torch
import torch.nn as nn
import numpy as np

class RunningMeanStd(nn.Module):
    def __init__(self, insize, epsilon=1e-05, per_channel=False, norm_only=False):
        super(RunningMeanStd, self).__init__()
        print('RunningMeanStd: ', insize)
        self.insize = insize
        self.epsilon = epsilon

        self.norm_only = norm_only
        self.per_channel = per_channel
        if per_channel:
            if len(self.insize) == 3:
                self.axis = [0,2,3]
            if len(self.insize) == 2:
                self.axis = [0,2]
            if len(self.insize) == 1:
                self.axis = [0]
            in_size = self.insize[0] 
        else:
            self.axis = [0]
            in_size = insize

        self.register_buffer('running_mean', torch.zeros(in_size, dtype = torch.float64))
        self.register_buffer('running_var', torch.ones(in_size, dtype = torch.float64))
        self.register_buffer('count', torch.ones((), dtype = torch.float64))

    def _update_mean_var_count_from_moments(self, mean, var, count, batch_mean, batch_var, batch_count):
        delta = batch_mean - mean
        tot_count = count + batch_count

        new_mean = mean + delta * batch_count / tot_count
        m_a = var * count
        m_b = batch_var * batch_count
        M2 = m_a + m_b + delta**2 * count * batch_count / tot_count
        new_var = M2 / tot_count
        new_count = tot_count
        return new_mean, new_var, new_count

    def forward(self, input, unnorm=False):
        if self.training:
            mean = input.mean(self.axis) # along channel axis
            var = input.var(self.axis)
            self.running_mean, self.running_var, self.count = self._update_mean_var_count_from_moments(self.running_mean, self.running_var, self.count, 
                                                    mean, var, input.size()[0] )

        # change shape
        if self.per_channel:
            if len(self.insize) == 3:
                current_mean = self.running_mean.view([1, self.insize[0], 1, 1]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0], 1, 1]).expand_as(input)
            if len(self.insize) == 2:
                current_mean = self.running_mean.view([1, self.insize[0], 1]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0], 1]).expand_as(input)
            if len(self.insize) == 1:
                current_mean = self.running_mean.view([1, self.insize[0]]).expand_as(input)
                current_var = self.running_var.view([1, self.insize[0]]).expand_as(input)        
        else:
            current_mean = self.running_mean
            current_var = self.running_var
        # get output


        if unnorm:
            y = torch.clamp(input, min=-5.0, max=5.0)
            y = torch.sqrt(current_var.float() + self.epsilon)*y + current_mean.float()
        else:
            if self.norm_only:
                y = input/ torch.sqrt(current_var.float() + self.epsilon)
            else:
                y = (input - current_mean.float()) / torch.sqrt(current_var.float() + self.epsilon)
                y = torch.clamp(y, min=-5.0, max=5.0)
        return y
NCF_policies/algo2/policy/models/running_mean_std.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# In-Hand Object Rotation via Rapid Motor Adaptation
# https://arxiv.org/abs/2210.04887
# Copyright (c) 2022 Haozhi Qi
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
# Based on: RLGames
# Copyright (c) 2019 Denys88
# Licence under MIT License
# https://github.com/Denys88/rl_games/
# --------------------------------------------------------

import gym
import torch
from torch.utils.data import Dataset


def transform_op(arr):
    """
    swap and then flatten axes 0 and 1
    """
    if arr is None:
        return arr
    s = arr.size()
    return arr.transpose(0, 1).reshape(s[0] * s[1], *s[2:])


class ExperienceBuffer(Dataset):
    def __init__(
        self,
        num_envs,
        horizon_length,
        batch_size,
        minibatch_size,
        obs_dim,
        act_dim,
        priv_dim,
        point_cloud_dim,
        device,
    ):
        self.device = device
        self.num_envs = num_envs
        self.transitions_per_env = horizon_length
        self.priv_info_dim = priv_dim

        self.data_dict = None
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.priv_dim = priv_dim
        self.point_cloud_dim = point_cloud_dim
        self.storage_dict = {
            "obses": torch.zeros(
                (self.transitions_per_env, self.num_envs, self.obs_dim),
                dtype=torch.float32,
                device=self.device,
            ),
            "priv_info": torch.zeros(
                (self.transitions_per_env, self.num_envs, self.priv_dim),
                dtype=torch.float32,
                device=self.device,
            ),
            "point_cloud_info": torch.zeros(
                (
                    self.transitions_per_env,
                    self.num_envs,
                    self.point_cloud_dim,
                    3,
                ),
                dtype=torch.float32,
                device=self.device,
            ),
            "rewards": torch.zeros(
                (self.transitions_per_env, self.num_envs, 1),
                dtype=torch.float32,
                device=self.device,
            ),
            "values": torch.zeros(
                (self.transitions_per_env, self.num_envs, 1),
                dtype=torch.float32,
                device=self.device,
            ),
            "neglogpacs": torch.zeros(
                (self.transitions_per_env, self.num_envs),
                dtype=torch.float32,
                device=self.device,
            ),
            "dones": torch.zeros(
                (self.transitions_per_env, self.num_envs),
                dtype=torch.uint8,
                device=self.device,
            ),
            "actions": torch.zeros(
                (self.transitions_per_env, self.num_envs, self.act_dim),
                dtype=torch.float32,
                device=self.device,
            ),
            "mus": torch.zeros(
                (self.transitions_per_env, self.num_envs, self.act_dim),
                dtype=torch.float32,
                device=self.device,
            ),
            "sigmas": torch.zeros(
                (self.transitions_per_env, self.num_envs, self.act_dim),
                dtype=torch.float32,
                device=self.device,
            ),
            "returns": torch.zeros(
                (self.transitions_per_env, self.num_envs, 1),
                dtype=torch.float32,
                device=self.device,
            ),
        }

        self.batch_size = batch_size
        self.minibatch_size = minibatch_size
        self.length = self.batch_size // self.minibatch_size

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        start = idx * self.minibatch_size
        end = (idx + 1) * self.minibatch_size
        self.last_range = (start, end)
        input_dict = {}
        for k, v in self.data_dict.items():
            if type(v) is dict:
                v_dict = {kd: vd[start:end] for kd, vd in v.items()}
                input_dict[k] = v_dict
            else:
                input_dict[k] = v[start:end]
        return (
            input_dict["values"],
            input_dict["neglogpacs"],
            input_dict["advantages"],
            input_dict["mus"],
            input_dict["sigmas"],
            input_dict["returns"],
            input_dict["actions"],
            input_dict["obses"],
            input_dict["priv_info"],
            input_dict["point_cloud_info"],
        )

    def update_mu_sigma(self, mu, sigma):
        start = self.last_range[0]
        end = self.last_range[1]
        self.data_dict["mus"][start:end] = mu
        self.data_dict["sigmas"][start:end] = sigma

    def update_data(self, name, index, val):
        if type(val) is dict:
            for k, v in val.items():
                self.storage_dict[name][k][index, :] = v
        else:
            self.storage_dict[name][index, :] = val

    def computer_return(self, last_values, gamma, tau):
        last_gae_lam = 0
        mb_advs = torch.zeros_like(self.storage_dict["rewards"])
        for t in reversed(range(self.transitions_per_env)):
            if t == self.transitions_per_env - 1:
                next_values = last_values
            else:
                next_values = self.storage_dict["values"][t + 1]
            next_nonterminal = 1.0 - self.storage_dict["dones"].float()[t]
            next_nonterminal = next_nonterminal.unsqueeze(1)
            delta = (
                self.storage_dict["rewards"][t]
                + gamma * next_values * next_nonterminal
                - self.storage_dict["values"][t]
            )
            mb_advs[t] = last_gae_lam = (
                delta + gamma * tau * next_nonterminal * last_gae_lam
            )
            self.storage_dict["returns"][t, :] = (
                mb_advs[t] + self.storage_dict["values"][t]
            )

    def prepare_training(self):
        self.data_dict = {}
        for k, v in self.storage_dict.items():
            self.data_dict[k] = transform_op(v)
        advantages = self.data_dict["returns"] - self.data_dict["values"]
        self.data_dict["advantages"] = (
            (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        ).squeeze(1)
        return self.data_dict

NCF_policies/algo2/policy/ppo/experience.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# In-Hand Object Rotation via Rapid Motor Adaptation
# https://arxiv.org/abs/2210.04887
# Copyright (c) 2022 Haozhi Qi
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
# Based on: RLGames
# Copyright (c) 2019 Denys88
# Licence under MIT License
# https://github.com/Denys88/rl_games/
# --------------------------------------------------------

import os
import time
import cv2
import torch
import torch.distributed as dist
import numpy as np

from algo2.policy.ppo.experience import ExperienceBuffer
from algo2.policy.models.models import ActorCritic
from algo2.policy.models.running_mean_std import RunningMeanStd

from algo2.utils.misc import AverageScalarMeter
from tensorboardX import SummaryWriter

import matplotlib.pyplot as plt
from algo2.utils.misc import add_to_fifo, multi_gpu_aggregate_stats
from algo2.ncf.utils.utils_ncf import NCF_dataloader, NCF_Viz


DEBUG = False


class PPO(object):
    def __init__(self, env, output_dir, full_config):
        current_file = os.path.abspath(__file__)
        self.root_dir = os.path.abspath(
            os.path.join(current_file, "..", "..", "..", "..", "..")
        )

        # ---- MultiGPU ----
        self.multi_gpu = full_config.train.ppo.multi_gpu
        if self.multi_gpu:
            self.rank = int(os.getenv("LOCAL_RANK", "0"))
            self.rank_size = int(os.getenv("WORLD_SIZE", "1"))
            dist.init_process_group("nccl", rank=self.rank, world_size=self.rank_size)
            self.device = "cuda:" + str(self.rank)
            print(f"current rank: {self.rank} and use device {self.device}")
        else:
            self.rank = -1
            self.device = full_config["rl_device"]
        self.network_config = full_config.train.network
        self.ppo_config = full_config.train.ppo

        # ---- build environment ----
        self.env = env
        self.num_actors = self.ppo_config["num_actors"]
        action_space = self.env.action_space
        self.actions_num = action_space.shape[0]
        self.actions_low = (
            torch.from_numpy(action_space.low.copy()).float().to(self.device)
        )
        self.actions_high = (
            torch.from_numpy(action_space.high.copy()).float().to(self.device)
        )
        self.observation_space = self.env.observation_space
        self.obs_shape = self.observation_space.shape

        # ---- Priv Info ----
        self.priv_info_dim = self.ppo_config["ncf_output_dim"]
        self.priv_info = self.ppo_config[
            "ncf_info"
        ]  # and self.ppo_config["ncf_use_gt"]
        # self.proprio_adapt = self.ppo_config["proprio_adapt"]

        # ---- Tactile Info ----
        self.tactile_info = self.ppo_config["tactile_info"]
        self.tactile_seq_length = self.ppo_config["tactile_seq_length"]
        self.tactile_info_embed_dim = self.ppo_config["tactile_info_embed_dim"]
        self.tactile_info_dim = (
            self.tactile_info_embed_dim * self.tactile_seq_length * 2
        )

        # ---- NCF Info ----
        self.ncf_info = self.ppo_config["ncf_info"]
        self.ncf_use_gt = self.ppo_config["ncf_use_gt"]
        self.ncf_output_dim = self.ppo_config["ncf_output_dim"]
        if self.ncf_info:
            self.tactile_info = False
        self.ncf_proprio_adapt = False

        # ---- Model ----
        net_config = {
            "actor_units": self.network_config.mlp.units,
            "actions_num": self.actions_num,
            "proprio_input_shape": self.obs_shape,
            "tactile_info": self.tactile_info,
            "ncf_info": self.ncf_info,
            "priv_info": self.priv_info,
            "tactile_units": self.network_config.tactile_mlp.units,
            "ncf_units": self.network_config.ncf_mlp.units,
            "ncf_adapt_units": self.network_config.ncf_adapt_mlp.units,
            "tactile_input_shape": (self.tactile_info_dim,),
            "ncf_input_shape": (self.ncf_output_dim,),
            "ncf_proprio_adapt": self.ncf_proprio_adapt,
        }
        self.model = ActorCritic(net_config)
        self.model.to(self.device)

        if self.tactile_info:
            self.obs_shape = (self.obs_shape[0] + self.tactile_info_dim,)

        self.obs_mean_std = RunningMeanStd(self.obs_shape).to(self.device)
        self.point_cloud_mean_std = RunningMeanStd(
            3,
        ).to(self.device)
        self.value_mean_std = RunningMeanStd((1,)).to(self.device)

        # ---- Output Dir ----
        # allows us to specify a folder where all experiments will reside
        self.output_dir = output_dir
        self.nn_dir = os.path.join(self.output_dir, "stage1_nn")
        self.tb_dif = os.path.join(self.output_dir, "stage1_tb")
        os.makedirs(self.nn_dir, exist_ok=True)
        os.makedirs(self.tb_dif, exist_ok=True)

        # ---- Optim ----
        self.last_lr = float(self.ppo_config["learning_rate"])
        self.weight_decay = self.ppo_config.get("weight_decay", 0.0)
        self.optimizer = torch.optim.Adam(
            self.model.parameters(), self.last_lr, weight_decay=self.weight_decay
        )

        # ---- PPO Train Param ----
        self.e_clip = self.ppo_config["e_clip"]
        self.clip_value = self.ppo_config["clip_value"]
        self.entropy_coef = self.ppo_config["entropy_coef"]
        self.critic_coef = self.ppo_config["critic_coef"]
        self.bounds_loss_coef = self.ppo_config["bounds_loss_coef"]
        self.gamma = self.ppo_config["gamma"]
        self.tau = self.ppo_config["tau"]
        self.truncate_grads = self.ppo_config["truncate_grads"]
        self.grad_norm = self.ppo_config["grad_norm"]
        self.value_bootstrap = self.ppo_config["value_bootstrap"]
        self.normalize_advantage = self.ppo_config["normalize_advantage"]
        self.normalize_input = self.ppo_config["normalize_input"]
        self.normalize_value = self.ppo_config["normalize_value"]
        self.normalize_point_cloud = self.ppo_config["normalize_point_cloud"]
        # ---- PPO Collect Param ----
        self.horizon_length = self.ppo_config["horizon_length"]
        self.batch_size = self.horizon_length * self.num_actors
        self.minibatch_size = self.ppo_config["minibatch_size"]
        self.mini_epochs_num = self.ppo_config["mini_epochs"]
        assert self.batch_size % self.minibatch_size == 0 or full_config.test
        # ---- scheduler ----
        self.kl_threshold = self.ppo_config["kl_threshold"]
        self.scheduler = AdaptiveScheduler(self.kl_threshold)
        # ---- Snapshot
        self.save_freq = self.ppo_config["save_frequency"]
        self.save_best_after = self.ppo_config["save_best_after"]
        # ---- Tensorboard Logger ----
        self.extra_info = {}
        writer = SummaryWriter(self.tb_dif)
        self.writer = writer

        # ---- Rollout GIFs ----
        self.gif_frame_counter = 0
        self.gif_save_every_n = 7500
        self.gif_save_length = 300
        self.gif_frames = []

        self.episode_rewards = AverageScalarMeter(50)
        self.episode_lengths = AverageScalarMeter(50)
        self.episode_success = AverageScalarMeter(50)
        self.obs = None
        self.epoch_num = 0

        self.storage = ExperienceBuffer(
            self.num_actors,
            self.horizon_length,
            self.batch_size,
            self.minibatch_size,
            self.obs_shape[0],
            self.actions_num,
            self.priv_info_dim,
            self.ncf_output_dim,
            self.device,
        )

        batch_size = self.num_actors
        current_rewards_shape = (batch_size, 1)
        self.current_rewards = torch.zeros(
            current_rewards_shape, dtype=torch.float32, device=self.device
        )
        self.current_lengths = torch.zeros(
            batch_size, dtype=torch.float32, device=self.device
        )
        self.dones = torch.ones((batch_size,), dtype=torch.uint8, device=self.device)
        self.agent_steps = 0
        self.max_agent_steps = self.ppo_config["max_agent_steps"]
        self.best_rewards = -10000

        # ---- Timing
        self.data_collect_time = 0
        self.rl_train_time = 0
        self.all_time = 0

        # ---- DIGIT autoencoder
        if self.tactile_info:
            path_checkpoint = full_config["task"]["ncf"]["checkpoint_vae"]
            self.digit_vae = load_digit_autoencoder(self.root_dir, path_checkpoint)
            self.digit_vae.to(self.device)
            self.digit_vae.eval()

        # ---- NCF
        if DEBUG:
            self.ncf_viz = NCF_Viz()

        if (self.ncf_info and not self.ncf_use_gt) or DEBUG:
            path_pointcloud_object = os.path.join(
                self.root_dir, full_config["task"]["ncf"]["path_pointcloud_object"]
            )
            path_ndf_code_object = os.path.join(
                self.root_dir, full_config["task"]["ncf"]["path_ndf_code_object"]
            )
            path_mesh_object = os.path.join(
                self.root_dir, full_config["task"]["ncf"]["path_mesh_object"]
            )
            self.ncf_dataloader = NCF_dataloader(
                num_envs=self.num_actors,
                path_mesh_object=path_mesh_object,
                path_pointcloud_object=path_pointcloud_object,
                path_ndf_code_object=path_ndf_code_object,
                pc_subsample=self.ppo_config["ncf_pc_subsample"],
            )
            path_checkpoint_vae = full_config["task"]["ncf"]["checkpoint_vae"]
            path_checkpoint_ndf = full_config["task"]["ncf"]["checkpoint_ndf"]

            path_checkpoint_ncf = "{0}/{1}_{2}/model-epoch={3:03d}.ckpt".format(
                self.root_dir,
                full_config["task"]["ncf"]["checkpoint_ncf"],
                full_config["task"]["ncf"]["ncf_arch"],
                full_config["task"]["ncf"]["ncf_epoch"],
            )

            self.digit_vae, self.ncf = load_ncf(
                full_config["task"]["ncf"]["ncf_arch"],
                self.root_dir,
                path_checkpoint_vae,
                path_checkpoint_ndf,
                path_checkpoint_ncf,
                self.device,
            )
            self.digit_vae.eval()
            self.digit_vae.to(self.device)
            self.ncf.eval()
            self.ncf.to(self.device)

        if (self.ncf_info or self.tactile_info) or DEBUG:
            self.reset_ncf_buffers()

    def reset_digits_fifo(self):
        self.digit_left_fifo = torch.zeros(
            (self.num_actors, self.tactile_seq_length, self.tactile_info_embed_dim),
            dtype=torch.float32,
            device=self.device,
        )
        self.digit_right_fifo = torch.zeros(
            (self.num_actors, self.tactile_seq_length, self.tactile_info_embed_dim),
            dtype=torch.float32,
            device=self.device,
        )

    def reset_eef_fifo(self):
        self.eef_fifo = torch.zeros(
            (self.num_actors, self.tactile_seq_length, 7),
            dtype=torch.float32,
            device=self.device,
        )
        self.eef_fifo[:, :, -1] = 1.0

    def reset_ncf_buffers(self):
        self.reset_digits_fifo()
        self.reset_eef_fifo()

    def write_stats(self, a_losses, c_losses, b_losses, entropies, kls, grad_norms):
        self.writer.add_scalar(
            "performance/RLTrainFPS",
            self.agent_steps / self.rl_train_time,
            self.agent_steps,
        )
        self.writer.add_scalar(
            "performance/EnvStepFPS",
            self.agent_steps / self.data_collect_time,
            self.agent_steps,
        )

        self.writer.add_scalar(
            "losses/actor_loss", torch.mean(a_losses).item(), self.agent_steps
        )
        self.writer.add_scalar(
            "losses/bounds_loss", torch.mean(b_losses).item(), self.agent_steps
        )
        self.writer.add_scalar(
            "losses/critic_loss", torch.mean(c_losses).item(), self.agent_steps
        )
        self.writer.add_scalar(
            "losses/entropy", torch.mean(entropies).item(), self.agent_steps
        )

        self.writer.add_scalar("info/last_lr", self.last_lr, self.agent_steps)
        self.writer.add_scalar("info/e_clip", self.e_clip, self.agent_steps)
        self.writer.add_scalar("info/kl", torch.mean(kls).item(), self.agent_steps)
        self.writer.add_scalar(
            "info/grad_norms", torch.mean(grad_norms).item(), self.agent_steps
        )

        for k, v in self.extra_info.items():
            self.writer.add_scalar(f"{k}", v, self.agent_steps)

    def set_eval(self):
        self.model.eval()
        if self.normalize_input:
            self.obs_mean_std.eval()
        if self.normalize_value:
            self.value_mean_std.eval()
        if self.normalize_point_cloud:
            self.point_cloud_mean_std.eval()

    def set_train(self):
        self.model.train()
        if self.normalize_input:
            self.obs_mean_std.train()
        if self.normalize_value:
            self.value_mean_std.train()
        if self.normalize_point_cloud:
            self.point_cloud_mean_std.train()

    def _get_ncf_output(self, obs_dict):
        if self.ncf_use_gt:
            contact_gt = obs_dict["gt_extrinsic_contact"]
            contact = contact_gt
            idx_query = np.arange(contact.shape[1])

        else:
            self._add_to_ncf_buffers(obs_dict)
            ee_pose, ee_t1, ee_t2 = self.ncf_dataloader.get_ee_sequence(self.eef_fifo)

            digits_emb_left = self.digit_left_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )
            digits_emb_right = self.digit_right_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )

            (
                shape_code,
                ncf_query_points,
                idx_query,
            ) = self.ncf_dataloader.get_pointcloud_data()

            inputs_ncf = {
                "digit_emb_left": digits_emb_left,
                "digit_emb_right": digits_emb_right,
                "ee_pose": ee_pose.to(self.device),
                "ee_pose_1": ee_t1.to(self.device),
                "ee_pose_2": ee_t2.to(self.device),
                "shape_emb": shape_code.to(self.device),
                "ncf_query_points": ncf_query_points.to(self.device),
            }
            with torch.no_grad():
                contact = self.ncf(inputs_ncf)

            d_mug_cupholder = (
                torch.norm(self.env.nut_pos - self.env.bolt_pos, dim=1)
            ) > 0.13
            contact[d_mug_cupholder] = 0.0
            contact = torch.where(contact > 0.4, 1.0, 0.0)
            contact_gt = contact.clone()

        if DEBUG:
            idx = 0
            object_pose = torch.cat(
                (self.env.nut_pos[idx], self.env.nut_quat[idx]), dim=0
            ).unsqueeze(0)
            cupholder_pose = torch.cat(
                (self.env.bolt_pos[idx], self.env.bolt_quat[idx]), dim=0
            ).unsqueeze(0)
            # p = contact[idx].cpu().numpy()
            p = np.zeros((contact[idx].shape[0]))
            p[idx_query] = contact[idx].cpu().numpy()
            # p[p > 0.3] = 1.0

            p_gt = contact_gt[idx].cpu().numpy()

            # img = self.ncf_viz.show_scene(
            #     object_pose, cupholder_pose, p_gt, p_pred, idx_query, interactive=False
            # )
            img_gt, img_pred = self.ncf_viz.show_gt_pred(
                object_pose,
                cupholder_pose,
                p_gt,
                p,
            )
            img = np.concatenate((img_gt, img_pred), axis=1)
            cv2.imshow("ncf", cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
            cv2.waitKey(1)

        return contact_gt, contact, idx_query

    def get_policy_inputs(self, obs_dict):
        if self.tactile_info:
            self._add_to_ncf_buffers(obs_dict)
            digits_emb_left = self.digit_left_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )
            digits_emb_right = self.digit_right_fifo.reshape(
                -1, self.tactile_seq_length * self.tactile_info_embed_dim
            )

            digits_emb = torch.cat([digits_emb_left, digits_emb_right], dim=-1)

            aug_obs = torch.cat([obs_dict["obs"], digits_emb], dim=-1)
            processed_obs = self.obs_mean_std(aug_obs)

        else:
            aug_obs = obs_dict["obs"]
            processed_obs = self.obs_mean_std(obs_dict["obs"])

        return processed_obs, aug_obs

    def model_act(self, obs_dict):
        processed_obs, aug_pos = self.get_policy_inputs(obs_dict)
        point_cloud = None

        if self.priv_info:
            ncf_priv_info, pred_contact, _ = self._get_ncf_output(obs_dict)
            point_cloud = obs_dict["pointclouds_t"]
            ncf_mask = torch.where(ncf_priv_info > 0.1, 1.0, 0.0)
            ncf_mask = ncf_mask.repeat(3, 1, 1).permute(1, 2, 0)
            point_cloud = point_cloud * ncf_mask
            obs_dict["pointclouds_t"] = point_cloud

            if self.normalize_point_cloud:
                point_cloud = self.point_cloud_mean_std(
                    point_cloud.reshape(-1, 3)
                ).reshape((processed_obs.shape[0], -1, 3))

        input_dict = {
            "obs": processed_obs,
            "priv_point_cloud": point_cloud,
        }

        res_dict = self.model.act(input_dict)
        res_dict["values"] = self.value_mean_std(res_dict["values"], True)
        return res_dict, aug_pos

    def train(self):
        _t = time.time()
        _last_t = time.time()
        self.obs = self.env.reset()
        self.reset_digits_fifo()
        self.agent_steps = (
            self.batch_size if not self.multi_gpu else self.batch_size * self.rank_size
        )

        if self.multi_gpu:
            torch.cuda.set_device(self.rank)
            print("====================broadcasting parameters")
            model_params = [self.model.state_dict()]
            dist.broadcast_object_list(model_params, 0)
            self.model.load_state_dict(model_params[0])

        while self.agent_steps < self.max_agent_steps:
            self.epoch_num += 1
            (
                a_losses,
                c_losses,
                b_losses,
                entropies,
                kls,
                grad_norms,
            ) = self.train_epoch()

            self.storage.data_dict = None

            (
                a_losses,
                b_losses,
                c_losses,
                entropies,
                kls,
                grad_norms,
            ) = multi_gpu_aggregate_stats(
                [a_losses, b_losses, c_losses, entropies, kls, grad_norms]
            )
            mean_rewards, mean_lengths, mean_success = multi_gpu_aggregate_stats(
                [
                    torch.Tensor([self.episode_rewards.get_mean()])
                    .float()
                    .to(self.device),
                    torch.Tensor([self.episode_lengths.get_mean()])
                    .float()
                    .to(self.device),
                    torch.Tensor([self.episode_success.get_mean()])
                    .float()
                    .to(self.device),
                ]
            )
            for k, v in self.extra_info.items():
                if type(v) is not torch.Tensor:
                    v = torch.Tensor([v]).float().to(self.device)
                self.extra_info[k] = multi_gpu_aggregate_stats(v[None].to(self.device))

            if not self.multi_gpu or (self.multi_gpu and self.rank == 0):
                all_fps = self.agent_steps / (time.time() - _t)
                last_fps = (
                    self.batch_size
                    if not self.multi_gpu
                    else self.batch_size * self.rank_size
                ) / (time.time() - _last_t)
                _last_t = time.time()
                info_string = (
                    f"Agent Steps: {int(self.agent_steps // 1e6):04}M | FPS: {all_fps:.1f} | "
                    f"Last FPS: {last_fps:.1f} | "
                    f"Collect Time: {self.data_collect_time / 60:.1f} min | "
                    f"Train RL Time: {self.rl_train_time / 60:.1f} min | "
                    f"Current Best: {self.best_rewards:.2f}"
                )
                print(info_string)

                self.write_stats(
                    a_losses, c_losses, b_losses, entropies, kls, grad_norms
                )
                self.writer.add_scalar(
                    "episode_rewards/step", mean_rewards, self.agent_steps
                )
                self.writer.add_scalar(
                    "episode_lengths/step", mean_lengths, self.agent_steps
                )
                self.writer.add_scalar(
                    "episode_success/step", mean_success, self.agent_steps
                )
                checkpoint_name = f"ep_{self.epoch_num}_step_{int(self.agent_steps // 1e6):04}m_reward_{mean_rewards:.2f}"

                if self.save_freq > 0:
                    if (self.epoch_num % self.save_freq == 0) and (
                        mean_rewards <= self.best_rewards
                    ):
                        self.save(os.path.join(self.nn_dir, checkpoint_name))
                        self.save(os.path.join(self.nn_dir, f"last"))

                if (
                    mean_rewards > self.best_rewards
                    and self.agent_steps >= self.save_best_after
                    and mean_rewards != 0.0
                ):
                    print(f"save current best reward: {mean_rewards:.2f}")
                    # remove previous best file
                    prev_best_ckpt = os.path.join(
                        self.nn_dir, f"best_reward_{self.best_rewards:.2f}.pth"
                    )
                    if os.path.exists(prev_best_ckpt):
                        os.remove(prev_best_ckpt)
                    self.best_rewards = mean_rewards
                    self.save(
                        os.path.join(self.nn_dir, f"best_reward_{mean_rewards:.2f}")
                    )

        print("max steps achieved")

    def save(self, name):
        weights = {
            "model": self.model.state_dict(),
        }
        if self.obs_mean_std:
            weights["running_mean_std"] = self.obs_mean_std.state_dict()
        if self.value_mean_std:
            weights["value_mean_std"] = self.value_mean_std.state_dict()
        if self.point_cloud_mean_std:
            weights["point_cloud_mean_std"] = self.point_cloud_mean_std.state_dict()
        torch.save(weights, f"{name}.pth")

    def restore_train(self, fn):
        if not fn:
            return
        checkpoint = torch.load(fn)
        for k in list(checkpoint["model"].keys()):
            print(k)

        self.model.load_state_dict(checkpoint["model"])
        self.obs_mean_std.load_state_dict(checkpoint["running_mean_std"])
        if self.normalize_point_cloud:
            self.point_cloud_mean_std.load_state_dict(
                checkpoint["point_cloud_mean_std"]
            )

    def restore_test(self, fn):
        checkpoint = torch.load(fn)
        self.model.load_state_dict(checkpoint["model"])
        if self.normalize_input:
            self.obs_mean_std.load_state_dict(checkpoint["running_mean_std"])
        if self.normalize_point_cloud:
            self.point_cloud_mean_std.load_state_dict(
                checkpoint["point_cloud_mean_std"]
            )
        self.played_games = 0

    def test(self):
        num_episodes = 5

        self.metrics = {}
        self.metrics["success"] = []
        self.metrics["steps"] = []
        self.metrics["dist_mug_cupholder"] = []
        self.metrics["error_aligment_keypoints"] = []
        self.metrics["error_quat"] = []
        self.metrics["mug_close_to_cupholder"] = []
        self.metrics["mug_oriented"] = []
        self.metrics["mean_time_complete_task"] = []
        imgs_env = []

        for n in range(num_episodes):
            self.set_eval()
            obs_dict = self.env.reset()
            self.reset_digits_fifo()

            while True:
                processed_obs, _ = self.get_policy_inputs(obs_dict)
                point_cloud = None

                if self.priv_info:
                    ncf_priv_info, pred_contact, _ = self._get_ncf_output(obs_dict)
                    point_cloud = obs_dict["pointclouds_t"]
                    ncf_mask = torch.where(ncf_priv_info > 0.1, 1.0, 0.0)
                    ncf_mask = ncf_mask.repeat(3, 1, 1).permute(1, 2, 0)
                    point_cloud = point_cloud * ncf_mask
                    obs_dict["pointclouds_t"] = point_cloud

                    if self.normalize_point_cloud:
                        point_cloud = self.point_cloud_mean_std(
                            point_cloud.reshape(-1, 3)
                        ).reshape((processed_obs.shape[0], -1, 3))

                input_dict = {
                    "obs": processed_obs,
                    "priv_point_cloud": point_cloud,
                }

                # if DEBUG:
                #     self._get_ncf_output(obs_dict)

                mu = self.model.act_inference(input_dict)
                mu = torch.clamp(mu, -1.0, 1.0)
                obs_dict, r, done, info = self.env.step(mu)

    def train_epoch(self):
        # collect minibatch data
        _t = time.time()
        self.set_eval()
        self.play_steps()
        self.data_collect_time += time.time() - _t
        # update network
        _t = time.time()
        self.set_train()
        a_losses, b_losses, c_losses = [], [], []
        entropies, kls, grad_norms = [], [], []
        for _ in range(0, self.mini_epochs_num):
            ep_kls = []
            for i in range(len(self.storage)):
                # value_preds, old_action_log_probs, advantage, old_mu, old_sigma, \
                #     returns, actions, obs, priv_info = self.storage[i]

                (
                    value_preds,
                    old_action_log_probs,
                    advantage,
                    old_mu,
                    old_sigma,
                    returns,
                    actions,
                    obs,
                    _,
                    point_cloud_info,
                ) = self.storage[i]

                obs = self.obs_mean_std(obs)
                if self.priv_info:
                    if self.normalize_point_cloud:
                        point_cloud_info = self.point_cloud_mean_std(
                            point_cloud_info.reshape(-1, 3)
                        ).reshape((obs.shape[0], -1, 3))

                batch_dict = {
                    "prev_actions": actions,
                    "obs": obs,
                    "priv_point_cloud": point_cloud_info,
                    # 'priv_info': priv_info,
                }
                res_dict = self.model(batch_dict)
                action_log_probs = res_dict["prev_neglogp"]
                values = res_dict["values"]
                entropy = res_dict["entropy"]
                mu = res_dict["mus"]
                sigma = res_dict["sigmas"]

                # actor loss
                ratio = torch.exp(old_action_log_probs - action_log_probs)
                surr1 = advantage * ratio
                surr2 = advantage * torch.clamp(
                    ratio, 1.0 - self.e_clip, 1.0 + self.e_clip
                )
                a_loss = torch.max(-surr1, -surr2)
                # critic loss
                value_pred_clipped = value_preds + (values - value_preds).clamp(
                    -self.e_clip, self.e_clip
                )
                value_losses = (values - returns) ** 2
                value_losses_clipped = (value_pred_clipped - returns) ** 2
                c_loss = torch.max(value_losses, value_losses_clipped)
                # bounded loss
                if self.bounds_loss_coef > 0:
                    soft_bound = 1.1
                    mu_loss_high = torch.clamp_min(mu - soft_bound, 0.0) ** 2
                    mu_loss_low = torch.clamp_max(mu + soft_bound, 0.0) ** 2
                    b_loss = (mu_loss_low + mu_loss_high).sum(axis=-1)
                else:
                    b_loss = torch.zeros_like(mu)
                a_loss, c_loss, entropy, b_loss = [
                    torch.mean(loss) for loss in [a_loss, c_loss, entropy, b_loss]
                ]

                loss = (
                    a_loss
                    + 0.5 * c_loss * self.critic_coef
                    - entropy * self.entropy_coef
                    + b_loss * self.bounds_loss_coef
                )

                self.optimizer.zero_grad()
                loss.backward()

                if self.multi_gpu:
                    # batch all_reduce ops: see https://github.com/entity-neural-network/incubator/pull/220
                    all_grads_list = []
                    for param in self.model.parameters():
                        if param.grad is not None:
                            all_grads_list.append(param.grad.view(-1))
                    all_grads = torch.cat(all_grads_list)
                    dist.all_reduce(all_grads, op=dist.ReduceOp.SUM)
                    offset = 0
                    for param in self.model.parameters():
                        if param.grad is not None:
                            param.grad.data.copy_(
                                all_grads[offset : offset + param.numel()].view_as(
                                    param.grad.data
                                )
                                / self.rank_size
                            )
                            offset += param.numel()

                grad_norms.append(
                    torch.norm(
                        torch.cat([p.reshape(-1) for p in self.model.parameters()])
                    )
                )

                if self.truncate_grads:
                    torch.nn.utils.clip_grad_norm_(
                        self.model.parameters(), self.grad_norm
                    )
                self.optimizer.step()

                with torch.no_grad():
                    kl_dist = policy_kl(mu.detach(), sigma.detach(), old_mu, old_sigma)

                kl = kl_dist
                a_losses.append(a_loss)
                c_losses.append(c_loss)
                ep_kls.append(kl)
                entropies.append(entropy)
                if self.bounds_loss_coef is not None:
                    b_losses.append(b_loss)

                self.storage.update_mu_sigma(mu.detach(), sigma.detach())

            av_kls = torch.mean(torch.stack(ep_kls))
            if self.multi_gpu:
                dist.all_reduce(av_kls, op=dist.ReduceOp.SUM)
                av_kls /= self.rank_size
            kls.append(av_kls)

            self.last_lr = self.scheduler.update(self.last_lr, av_kls.item())
            if self.multi_gpu:
                lr_tensor = torch.tensor([self.last_lr], device=self.device)
                dist.broadcast(lr_tensor, 0)
                lr = lr_tensor.item()
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = lr
            else:
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = self.last_lr

        self.rl_train_time += time.time() - _t
        return a_losses, c_losses, b_losses, entropies, kls, grad_norms

    def play_steps(self):
        for n in range(self.horizon_length):
            res_dict, aug_pos = self.model_act(self.obs)

            # collect o_t
            # self.storage.update_data("obses", n, self.obs["obs"])
            self.storage.update_data("obses", n, aug_pos)
            if self.ncf_info:
                self.storage.update_data(
                    "point_cloud_info", n, self.obs["pointclouds_t"]
                )
            # self.storage.update_data('priv_info', n, self.obs['priv_info'])
            for k in ["actions", "neglogpacs", "values", "mus", "sigmas"]:
                self.storage.update_data(k, n, res_dict[k])
            # do env step
            actions = torch.clamp(res_dict["actions"], -1.0, 1.0)

            # render() is called during env.step()
            # to save time, save gif only per gif_save_every_n steps
            # 1 step = #gpu * #envs agent steps
            record_frame = False
            # if self.gif_frame_counter % self.gif_save_every_n < self.gif_save_length:
            #     record_frame = True
            # record_frame = record_frame and int(os.getenv('LOCAL_RANK', '0')) == 0
            # self.env.enable_camera_sensors = record_frame
            # self.gif_frame_counter += 1

            self.obs, rewards, self.dones, infos = self.env.step(actions)

            if record_frame:
                self.gif_frames.append(self.env.capture_frame())
                # add frame to GIF
                if len(self.gif_frames) == self.gif_save_length:
                    frame_array = np.array([f["color"] for f in self.gif_frames])[
                        None
                    ]  # add batch axis
                    self.writer.add_video(
                        "rollout_gif",
                        frame_array,
                        global_step=self.agent_steps,
                        dataformats="NTHWC",
                        fps=20,
                    )
                    self.writer.flush()
                    self.gif_frames.clear()

            rewards = rewards.unsqueeze(1)
            # update dones and rewards after env step
            self.storage.update_data("dones", n, self.dones)
            shaped_rewards = 0.01 * rewards.clone()
            if self.value_bootstrap and "time_outs" in infos:
                shaped_rewards += (
                    self.gamma
                    * res_dict["values"]
                    * infos["time_outs"].unsqueeze(1).float()
                )
            self.storage.update_data("rewards", n, shaped_rewards)

            self.current_rewards += rewards
            self.current_lengths += 1
            done_indices = self.dones.nonzero(as_tuple=False)
            self.episode_rewards.update(self.current_rewards[done_indices])
            self.episode_lengths.update(self.current_lengths[done_indices])

            assert isinstance(infos, dict), "Info Should be a Dict"
            self.extra_info = {}
            for k, v in infos.items():
                # only log scalars
                if (
                    isinstance(v, float)
                    or isinstance(v, int)
                    or (isinstance(v, torch.Tensor) and len(v.shape) == 0)
                ):
                    self.extra_info[k] = v

            not_dones = 1.0 - self.dones.float()

            self.current_rewards = self.current_rewards * not_dones.unsqueeze(1)
            self.current_lengths = self.current_lengths * not_dones

        res_dict, aug_pos = self.model_act(self.obs)
        last_values = res_dict["values"]

        self.agent_steps = (
            (self.agent_steps + self.batch_size)
            if not self.multi_gpu
            else self.agent_steps + self.batch_size * self.rank_size
        )
        self.storage.computer_return(last_values, self.gamma, self.tau)
        self.storage.prepare_training()

        returns = self.storage.data_dict["returns"]
        values = self.storage.data_dict["values"]
        if self.normalize_value:
            self.value_mean_std.train()
            values = self.value_mean_std(values)
            returns = self.value_mean_std(returns)
            self.value_mean_std.eval()
        self.storage.data_dict["values"] = values
        self.storage.data_dict["returns"] = returns

    def _digit_encode(self, images):
        images = images.permute(0, 1, 4, 2, 3)
        with torch.no_grad():
            images_hat, digit_embeddings = self.digit_vae(images)

        return digit_embeddings.unsqueeze(1), images_hat

    def _add_to_ncf_buffers(self, obs_dict):
        if obs_dict["obs"].sum() != 0:
            if self.eef_fifo[:, :, 0:3].sum() == 0:
                self.eef_fifo = (
                    obs_dict["obs"][:, 0:7]
                    .unsqueeze(1)
                    .repeat(1, self.tactile_seq_length, 1)
                )
            else:
                self.eef_fifo = add_to_fifo(
                    self.eef_fifo, obs_dict["obs"][:, 0:7].unsqueeze(1)
                )

        digits_imgs_left = obs_dict["digits_left"]
        digits_imgs_right = obs_dict["digits_right"]

        digits_emb_left, left_hat = self._digit_encode(digits_imgs_left)
        digits_emb_right, right_hat = self._digit_encode(digits_imgs_right)

        self.digit_left_fifo = add_to_fifo(self.digit_left_fifo, digits_emb_left)
        self.digit_right_fifo = add_to_fifo(self.digit_right_fifo, digits_emb_right)


def policy_kl(p0_mu, p0_sigma, p1_mu, p1_sigma):
    c1 = torch.log(p1_sigma / p0_sigma + 1e-5)
    c2 = (p0_sigma**2 + (p1_mu - p0_mu) ** 2) / (2.0 * (p1_sigma**2 + 1e-5))
    c3 = -1.0 / 2.0
    kl = c1 + c2 + c3
    kl = kl.sum(dim=-1)  # returning mean between all steps of sum between all actions
    return kl.mean()


def load_digit_autoencoder(root_dir, path_checkpoint):
    from algo2.ncf.vae.vae import VAE

    path_checkpoint = os.path.join(root_dir, path_checkpoint)
    vae = VAE.load_from_checkpoint(
        path_checkpoint,
        enc_out_dim=512,
        latent_dim=64,
        input_height=64,
    )
    vae.eval()
    for param in vae.parameters():
        param.requires_grad = False
    return vae


def load_ndf(root_dir, path_checkpoint):
    from algo2.ncf.ndf.ndf import NDF

    path_checkpoint = os.path.join(root_dir, path_checkpoint)
    ndf = NDF(latent_dim=256, return_features=True, sigmoid=True)
    checkpoint = torch.load(path_checkpoint)
    ndf_weights = ndf.state_dict()

    for key in ndf_weights.keys():
        ndf_weights[key] = checkpoint[key]
    ndf.load_state_dict(ndf_weights)
    for param in ndf.parameters():
        param.requires_grad = False

    return ndf


def load_ncf(
    arch,
    root_dir,
    path_checkpoint_vae,
    path_checkpoint_ndf,
    path_checkpoint_ncf,
    device_sim,
):
    from algo2.ncf.contact_regressor.ncf_mlp import NCF as NCF_mlp
    from algo2.ncf.contact_regressor.ncf_transformer import NCF as NCF_transformer
    from algo2.ncf.config.config import NCF_Params
    from algo2.ncf.pipeline import NCF_Pipeline

    cfg = NCF_Params()
    digit_encoder = load_digit_autoencoder(root_dir, path_checkpoint_vae)
    ndf_model = load_ndf(root_dir, path_checkpoint_ndf)

    if "mlp" in arch:
        ncf_model = NCF_mlp(cfg)
    elif "transformer" in arch:
        ncf_model = NCF_transformer(cfg)
    else:
        raise NotImplementedError

    pipeline = NCF_Pipeline.load_from_checkpoint(
        path_checkpoint_ncf,
        cfg=cfg,
        digit_vae=digit_encoder,
        ndf=ndf_model,
        ncf=ncf_model,
    )

    for param in pipeline.parameters():
        param.requires_grad = False

    return pipeline.digit_vae, pipeline.ncf

    # checkpoint = torch.load(path_checkpoint_ncf)
    # ncf_weights = ncf_model.state_dict()

    # for key in ncf_weights.keys():
    #     ncf_weights[key] = checkpoint["state_dict"]["ncf." + key]
    # ncf_model.load_state_dict(ncf_weights)
    # for param in ncf_model.parameters():
    #     param.requires_grad = False

    # return digit_encoder, ncf_model


# def load_ncf_old(
#     arch,
#     root_dir,
#     path_checkpoint_vae,
#     path_checkpoint_ndf,
#     path_checkpoint_ncf,
#     device_sim,
#     device_ndf,
#     device_ncf,
# ):
#     from algo.ncf.ncf.ncf_mlp import NCF as NCF_mlp
#     from algo.ncf.ncf.ncf_transformer import NCF as NCF_transformer
#     from algo.ncf.pipeline import NCF_Pipeline
#     from algo.ncf.config.config import NCF_Params

#     cfg = NCF_Params(ncf_arch=arch)

#     digit_encoder = load_digit_autoencoder(root_dir, path_checkpoint_vae).to(device_sim)
#     ndf_model = load_ndf(root_dir, path_checkpoint_ndf, device_ndf)
#     # ndf_model = torch.nn.DataParallel(ndf_model, device_ids=[1, 2, 3])
#     ndf_model.to(device_ndf)

#     if arch == "mlp":
#         ncf_model = NCF_mlp(cfg)
#     elif arch == "transformer":
#         ncf_model = NCF_transformer(cfg)
#     else:
#         raise NotImplementedError

#     checkpoint_ncf = os.path.join(root_dir, path_checkpoint_ncf)
#     ncf_model.load_state_dict(torch.load(checkpoint_ncf, map_location=device_ncf))
#     if not DEBUG:
#         ncf_model = torch.nn.DataParallel(ncf_model, device_ids=[2, 3])
#     ncf_model.to(device_ncf)
#     ncf_model.eval()
#     for param in ncf_model.parameters():
#         param.requires_grad = False
#     return digit_encoder, ndf_model, ncf_model

#     # ncf_model.to(device_ncf)
#     # checkpoint_ncf = os.path.join(root_dir, path_checkpoint_ncf)
#     # ncf_pipeline = NCF_Pipeline.load_from_checkpoint(
#     #     checkpoint_ncf,
#     #     cfg=cfg,
#     #     digit_vae=digit_encoder,
#     #     ndf=ndf_model,
#     #     ncf=ncf_model,
#     # )
#     # ncf_pipeline.eval()
#     # ncf_pipeline.to(device_ncf)
#     # for param in ncf_pipeline.parameters():
#     #     param.requires_grad = False

#     # return ncf_pipeline


# from https://github.com/leggedrobotics/rsl_rl/blob/master/rsl_rl/algorithms/ppo.py
class AdaptiveScheduler(object):
    def __init__(self, kl_threshold=0.008):
        super().__init__()
        self.min_lr = 1e-6
        self.max_lr = 1e-2
        self.kl_threshold = kl_threshold

    def update(self, current_lr, kl_dist):
        lr = current_lr
        if kl_dist > (2.0 * self.kl_threshold):
            lr = max(current_lr / 1.5, self.min_lr)
        if kl_dist < (0.5 * self.kl_threshold):
            lr = min(current_lr * 1.5, self.max_lr)
        return lr

NCF_policies/algo2/policy/ppo/ppo.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# In-Hand Object Rotation via Rapid Motor Adaptation
# https://arxiv.org/abs/2210.04887
# Copyright (c) 2022 Haozhi Qi
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
# Based on: IsaacGymEnvs
# Copyright (c) 2018-2022, NVIDIA Corporation
# Licence under BSD 3-Clause License
# https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/
# --------------------------------------------------------

import os
import torch
import shlex
import random
import subprocess
import numpy as np
import torch.distributed as dist


def tprint(*args):
    """Temporarily prints things on the screen"""
    print("\r", end="")
    print(*args, end="")


def pprint(*args):
    """Permanently prints things on the screen"""
    print("\r", end="")
    print(*args)


def git_hash():
    cmd = 'git log -n 1 --pretty="%h"'
    ret = subprocess.check_output(shlex.split(cmd)).strip()
    if isinstance(ret, bytes):
        ret = ret.decode()
    return ret


def git_diff_config(name):
    cmd = f"git diff --unified=0 {name}"
    ret = subprocess.check_output(shlex.split(cmd)).strip()
    if isinstance(ret, bytes):
        ret = ret.decode()
    return ret


def set_np_formatting():
    """formats numpy print"""
    np.set_printoptions(
        edgeitems=30,
        infstr="inf",
        linewidth=4000,
        nanstr="nan",
        precision=2,
        suppress=False,
        threshold=10000,
        formatter=None,
    )


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

    return seed


def get_rank():
    return int(os.getenv("LOCAL_RANK", "0"))


def get_world_size():
    # return number of gpus
    if "LOCAL_WORLD_SIZE" in os.environ.keys():
        world_size = int(os.environ["LOCAL_WORLD_SIZE"])
    else:
        world_size = 1
    return world_size


def multi_gpu_aggregate_stats(values):
    if type(values) is not list:
        single_item = True
        values = [values]
    else:
        single_item = False
    rst = []
    for v in values:
        if type(v) is list:
            v = torch.stack(v)
        if get_world_size() > 1:
            dist.all_reduce(v, op=dist.ReduceOp.SUM)
            v = v / get_world_size()
        if v.numel() == 1:
            v = v.item()
        rst.append(v)
    if single_item:
        rst = rst[0]
    return rst


def add_to_fifo(tensor, x):
    """Pushes a new value to a tensor, removing the oldest one"""
    # return torch.cat((tensor[:,1:], x), dim=1)
    return torch.cat((x, tensor[:, 0:-1]), dim=1)


def multi_gpu_aggregate_stats(values):
    # lazy imports to avoid isaacgym errors
    import torch
    import torch.distributed as dist

    if type(values) is not list:
        single_item = True
        values = [values]
    else:
        single_item = False
    rst = []
    for v in values:
        if type(v) is list:
            v = torch.stack(v)
        if get_world_size() > 1:
            dist.all_reduce(v, op=dist.ReduceOp.SUM)
            v = v / get_world_size()
        if v.numel() == 1:
            v = v.item()
        rst.append(v)
    if single_item:
        rst = rst[0]
    return rst


class AverageScalarMeter(object):
    def __init__(self, window_size):
        self.window_size = window_size
        self.current_size = 0
        self.mean = 0

    def update(self, values):
        size = values.size()[0]
        if size == 0:
            return
        new_mean = torch.mean(values.float(), dim=0).cpu().numpy().item()
        size = np.clip(size, 0, self.window_size)
        old_size = min(self.window_size - size, self.current_size)
        size_sum = old_size + size
        self.current_size = size_sum
        self.mean = (self.mean * old_size + new_mean * size) / size_sum

    def clear(self):
        self.current_size = 0
        self.mean = 0

    def __len__(self):
        return self.current_size

    def get_mean(self):
        return self.mean

NCF_policies/algo2/utils/misc.py
carolinahiguera/NCF-Policy
import hydra
from hydra import compose, initialize
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, OmegaConf
from isaacgymenvs.utils.reformat import omegaconf_to_dict


OmegaConf.register_new_resolver('eq', lambda x, y: x.lower()==y.lower())
OmegaConf.register_new_resolver('contains', lambda x, y: x.lower() in y.lower())
OmegaConf.register_new_resolver('if', lambda pred, a, b: a if pred else b)
OmegaConf.register_new_resolver('resolve_default', lambda default, arg: default if arg=='' else arg)


# def make(
#     seed: int, 
#     task: str, 
#     num_envs: int, 
#     sim_device: str,
#     rl_device: str,
#     graphics_device_id: int = -1,
#     headless: bool = False,
#     multi_gpu: bool = False,
#     virtual_screen_capture: bool = False,
#     force_render: bool = True,
#     cfg: DictConfig = None
# ): 
#     from isaacgymenvs.utils.rlgames_utils import get_rlgames_env_creator
#     # create hydra config if no config passed in
#     if cfg is None:
#         # reset current hydra config if already parsed (but not passed in here)
#         if HydraConfig.initialized():
#             task = HydraConfig.get().runtime.choices['task']
#             hydra.core.global_hydra.GlobalHydra.instance().clear()

#         with initialize(config_path="./cfg"):
#             cfg = compose(config_name="config", overrides=[f"task={task}"])
#             cfg_dict = omegaconf_to_dict(cfg.task)
#             cfg_dict['env']['numEnvs'] = num_envs
#     # reuse existing config
#     else:
#         cfg_dict = omegaconf_to_dict(cfg.task)

#     create_rlgpu_env = get_rlgames_env_creator(
#         seed=seed,
#         task_config=cfg_dict,
#         task_name=cfg_dict["name"],
#         sim_device=sim_device,
#         rl_device=rl_device,
#         graphics_device_id=graphics_device_id,
#         headless=headless,
#         multi_gpu=multi_gpu,
#         virtual_screen_capture=virtual_screen_capture,
#         force_render=force_render,
#     )
#     return create_rlgpu_env()

NCF_policies/isaacgymenvs/__init__.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from .ncf_tacto.factory_task_mug_cupholder import NCFTaskCupholder as NCFTaskCupholder
from .ncf_tacto.factory_task_bowl_dishrack import NCFTaskDishrack as NCFTaskDishrack

# Mappings from strings to environments
isaacgym_task_map = {
    "NCFTaskMugCupholder": NCFTaskCupholder,
    "NCFTaskBowlDishrack": NCFTaskDishrack,
}

NCF_policies/isaacgymenvs/tasks/__init__.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import os
import time
from datetime import datetime
from os.path import join
from typing import Dict, Any, Tuple, List, Set

import gym
from gym import spaces

from isaacgym import gymtorch, gymapi
from isaacgymenvs.utils.torch_jit_utils import to_torch
from isaacgymenvs.utils.dr_utils import (
    get_property_setter_map,
    get_property_getter_map,
    get_default_setter_args,
    apply_random_samples,
    check_buckets,
    generate_random_samples,
)

import torch
import numpy as np
import operator, random
from copy import deepcopy
from isaacgymenvs.utils.utils import nested_dict_get_attr, nested_dict_set_attr

from collections import deque

import sys

import abc
from abc import ABC

EXISTING_SIM = None
SCREEN_CAPTURE_RESOLUTION = (1027, 768)


def _create_sim_once(gym, *args, **kwargs):
    global EXISTING_SIM
    if EXISTING_SIM is not None:
        return EXISTING_SIM
    else:
        EXISTING_SIM = gym.create_sim(*args, **kwargs)
        return EXISTING_SIM


class Env(ABC):
    def __init__(
        self,
        config: Dict[str, Any],
        rl_device: str,
        sim_device: str,
        graphics_device_id: int,
        headless: bool,
    ):
        """Initialise the env.

        Args:
            config: the configuration dictionary.
            sim_device: the device to simulate physics on. eg. 'cuda:0' or 'cpu'
            graphics_device_id: the device ID to render with.
            headless: Set to False to disable viewer rendering.
        """

        split_device = sim_device.split(":")
        self.device_type = split_device[0]
        self.device_id = int(split_device[1]) if len(split_device) > 1 else 0

        self.device = "cpu"
        if config["sim"]["use_gpu_pipeline"]:
            if self.device_type.lower() == "cuda" or self.device_type.lower() == "gpu":
                self.device = "cuda" + ":" + str(self.device_id)
            else:
                print(
                    "GPU Pipeline can only be used with GPU simulation. Forcing CPU Pipeline."
                )
                config["sim"]["use_gpu_pipeline"] = False

        self.rl_device = rl_device

        # Rendering
        # if training in a headless mode
        self.headless = headless

        enable_camera_sensors = config.get("enableCameraSensors", False)
        self.graphics_device_id = graphics_device_id
        if enable_camera_sensors == False and self.headless == True:
            self.graphics_device_id = -1

        self.num_environments = config["env"]["numEnvs"]
        self.num_agents = config["env"].get(
            "numAgents", 1
        )  # used for multi-agent environments

        self.num_observations = config["env"].get("numObservations", 0)
        self.num_states = config["env"].get("numStates", 0)

        self.obs_space = spaces.Box(
            np.ones(self.num_obs) * -np.Inf, np.ones(self.num_obs) * np.Inf
        )
        self.state_space = spaces.Box(
            np.ones(self.num_states) * -np.Inf, np.ones(self.num_states) * np.Inf
        )

        self.num_actions = config["env"]["numActions"]
        self.control_freq_inv = config["env"].get("controlFrequencyInv", 1)

        self.act_space = spaces.Box(
            np.ones(self.num_actions) * -1.0, np.ones(self.num_actions) * 1.0
        )

        self.clip_obs = config["env"].get("clipObservations", np.Inf)
        self.clip_actions = config["env"].get("clipActions", np.Inf)

        # Total number of training frames since the beginning of the experiment.
        # We get this information from the learning algorithm rather than tracking ourselves.
        # The learning algorithm tracks the total number of frames since the beginning of training and accounts for
        # experiments restart/resumes. This means this number can be > 0 right after initialization if we resume the
        # experiment.
        self.total_train_env_frames: int = 0

        # number of control steps
        self.control_steps: int = 0

        self.render_fps: int = config["env"].get("renderFPS", -1)
        self.last_frame_time: float = 0.0

        self.record_frames: bool = False
        self.record_frames_dir = join(
            "recorded_frames", datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        )

    @abc.abstractmethod
    def allocate_buffers(self):
        """Create torch buffers for observations, rewards, actions dones and any additional data."""

    @abc.abstractmethod
    def step(
        self, actions: torch.Tensor
    ) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, Dict[str, Any]]:
        """Step the physics of the environment.
        Args:
            actions: actions to apply
        Returns:
            Observations, rewards, resets, info
            Observations are dict of observations (currently only one member called 'obs')
        """

    @abc.abstractmethod
    def reset(self) -> Dict[str, torch.Tensor]:
        """Reset the environment.
        Returns:
            Observation dictionary
        """

    @abc.abstractmethod
    def reset_idx(self, env_ids: torch.Tensor):
        """Reset environments having the provided indices.
        Args:
            env_ids: environments to reset
        """

    @property
    def observation_space(self) -> gym.Space:
        """Get the environment's observation space."""
        return self.obs_space

    @property
    def action_space(self) -> gym.Space:
        """Get the environment's action space."""
        return self.act_space

    @property
    def num_envs(self) -> int:
        """Get the number of environments."""
        return self.num_environments

    @property
    def num_acts(self) -> int:
        """Get the number of actions in the environment."""
        return self.num_actions

    @property
    def num_obs(self) -> int:
        """Get the number of observations in the environment."""
        return self.num_observations

    def set_train_info(self, env_frames, *args, **kwargs):
        """
        Send the information in the direction algo->environment.
        Most common use case: tell the environment how far along we are in the training process. This is useful
        for implementing curriculums and things such as that.
        """
        self.total_train_env_frames = env_frames
        # print(f'env_frames updated to {self.total_train_env_frames}')

    def get_env_state(self):
        """
        Return serializable environment state to be saved to checkpoint.
        Can be used for stateful training sessions, i.e. with adaptive curriculums.
        """
        return None

    def set_env_state(self, env_state):
        pass


class VecTask(Env):
    metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 24}

    def __init__(
        self,
        config,
        rl_device,
        sim_device,
        graphics_device_id,
        headless,
        virtual_screen_capture: bool = False,
        force_render: bool = False,
    ):
        """Initialise the `VecTask`.

        Args:
            config: config dictionary for the environment.
            sim_device: the device to simulate physics on. eg. 'cuda:0' or 'cpu'
            graphics_device_id: the device ID to render with.
            headless: Set to False to disable viewer rendering.
            virtual_screen_capture: Set to True to allow the users get captured screen in RGB array via `env.render(mode='rgb_array')`.
            force_render: Set to True to always force rendering in the steps (if the `control_freq_inv` is greater than 1 we suggest stting this arg to True)
        """
        # super().__init__(config, rl_device, sim_device, graphics_device_id, headless, use_dict_obs)
        super().__init__(config, rl_device, sim_device, graphics_device_id, headless)
        self.virtual_screen_capture = virtual_screen_capture
        self.virtual_display = None
        if self.virtual_screen_capture:
            from pyvirtualdisplay.smartdisplay import SmartDisplay

            self.virtual_display = SmartDisplay(size=SCREEN_CAPTURE_RESOLUTION)
            self.virtual_display.start()
        self.force_render = force_render

        self.sim_params = self.__parse_sim_params(
            self.cfg["physics_engine"], self.cfg["sim"]
        )
        if self.cfg["physics_engine"] == "physx":
            self.physics_engine = gymapi.SIM_PHYSX
        elif self.cfg["physics_engine"] == "flex":
            self.physics_engine = gymapi.SIM_FLEX
        else:
            msg = f"Invalid physics engine backend: {self.cfg['physics_engine']}"
            raise ValueError(msg)

        self.dt: float = self.sim_params.dt

        # optimization flags for pytorch JIT
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)

        self.gym = gymapi.acquire_gym()

        self.first_randomization = True
        self.original_props = {}
        self.dr_randomizations = {}
        self.actor_params_generator = None
        self.extern_actor_params = {}
        self.last_step = -1
        self.last_rand_step = -1
        for env_id in range(self.num_envs):
            self.extern_actor_params[env_id] = None

        # create envs, sim and viewer
        self.sim_initialized = False
        self.create_sim()
        self.gym.prepare_sim(self.sim)
        self.sim_initialized = True

        self.set_viewer()
        self.allocate_buffers()

        self.obs_dict = {}

        self.gym.set_light_parameters(
            self.sim,
            0,
            gymapi.Vec3(0.8, 0.8, 0.8),
            gymapi.Vec3(0.8, 0.8, 0.8),
            gymapi.Vec3(-1, -2, 3),
        )

    def set_viewer(self):
        """Create the viewer."""

        # todo: read from config
        self.enable_viewer_sync = True
        self.viewer = None

        # if running with a viewer, set up keyboard shortcuts and camera
        if self.headless == False:
            # subscribe to keyboard shortcuts
            self.viewer = self.gym.create_viewer(self.sim, gymapi.CameraProperties())
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_ESCAPE, "QUIT"
            )
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_V, "toggle_viewer_sync"
            )
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_R, "record_frames"
            )

            # set the camera position based on up axis
            sim_params = self.gym.get_sim_params(self.sim)
            if sim_params.up_axis == gymapi.UP_AXIS_Z:
                cam_pos = gymapi.Vec3(20.0, 25.0, 3.0)
                cam_target = gymapi.Vec3(10.0, 15.0, 0.0)
            else:
                cam_pos = gymapi.Vec3(20.0, 3.0, 25.0)
                cam_target = gymapi.Vec3(10.0, 0.0, 15.0)

            self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

    def allocate_buffers(self):
        """Allocate the observation, states, etc. buffers.

        These are what is used to set observations and states in the environment classes which
        inherit from this one, and are read in `step` and other related functions.

        """

        # allocate buffers
        self.obs_buf = torch.zeros(
            (self.num_envs, self.num_obs), device=self.device, dtype=torch.float
        )
        self.states_buf = torch.zeros(
            (self.num_envs, self.num_states), device=self.device, dtype=torch.float
        )
        self.rew_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.float)
        self.reset_buf = torch.ones(self.num_envs, device=self.device, dtype=torch.long)
        self.success_buf = torch.ones(
            self.num_envs, device=self.device, dtype=torch.long
        )
        self.timeout_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long
        )
        self.progress_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long
        )
        self.randomize_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long
        )
        self.extras = {}

        # ncf buffers
        digit_sz = self.cfg["ncf"]["digits_size"]
        digit_seq_len = self.cfg["ncf"]["digits_seq_length"]
        self.digit_left_buf = torch.zeros(
            (self.num_envs, digit_seq_len, digit_sz, digit_sz, 3),
            device=self.device,
            dtype=torch.float,
        )
        self.digit_right_buf = torch.zeros(
            (self.num_envs, digit_seq_len, digit_sz, digit_sz, 3),
            device=self.device,
            dtype=torch.float,
        )

        self.gt_extrinsic_contact = torch.zeros(
            (self.num_envs, 2000), device=self.device, dtype=torch.float
        )

        self.pointclouds_t = torch.zeros(
            (self.num_envs, 2000, 3), device=self.device, dtype=torch.float
        )

        # self.digit_left_buf = np.zeros((self.num_envs, digit_sz, digit_sz, 3), dtype=np.float32)
        # self.digit_right_buf = np.zeros((self.num_envs, digit_sz, digit_sz, 3), dtype=np.float32)

    def create_sim(
        self,
        compute_device: int,
        graphics_device: int,
        physics_engine,
        sim_params: gymapi.SimParams,
    ):
        """Create an Isaac Gym sim object.

        Args:
            compute_device: ID of compute device to use.
            graphics_device: ID of graphics device to use.
            physics_engine: physics engine to use (`gymapi.SIM_PHYSX` or `gymapi.SIM_FLEX`)
            sim_params: sim params to use.
        Returns:
            the Isaac Gym sim object.
        """
        sim = _create_sim_once(
            self.gym, compute_device, graphics_device, physics_engine, sim_params
        )
        if sim is None:
            print("*** Failed to create sim")
            quit()

        return sim

    def get_state(self):
        """Returns the state buffer of the environment (the privileged observations for asymmetric training)."""
        return torch.clamp(self.states_buf, -self.clip_obs, self.clip_obs).to(
            self.rl_device
        )

    @abc.abstractmethod
    def pre_physics_step(self, actions: torch.Tensor):
        """Apply the actions to the environment (eg by setting torques, position targets).

        Args:
            actions: the actions to apply
        """

    @abc.abstractmethod
    def post_physics_step(self):
        """Compute reward and observations, reset any environments that require it."""

    def step(
        self, actions: torch.Tensor
    ) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, Dict[str, Any]]:
        """Step the physics of the environment.

        Args:
            actions: actions to apply
        Returns:
            Observations, rewards, resets, info
            Observations are dict of observations (currently only one member called 'obs')
        """

        # randomize actions
        if self.dr_randomizations.get("actions", None):
            actions = self.dr_randomizations["actions"]["noise_lambda"](actions)

        action_tensor = torch.clamp(actions, -self.clip_actions, self.clip_actions)
        # apply actions
        self.pre_physics_step(action_tensor)

        # step physics and render each frame
        for i in range(self.control_freq_inv):
            if self.force_render:
                self.render()
            self.gym.simulate(self.sim)

        # to fix!
        if self.device == "cpu":
            self.gym.fetch_results(self.sim, True)

        # compute observations, rewards, resets, ...
        self.post_physics_step()

        self.control_steps += 1

        # fill time out buffer: set to 1 if we reached the max episode length AND the reset buffer is 1. Timeout == 1 makes sense only if the reset buffer is 1.
        # self.timeout_buf = (self.progress_buf >= self.max_episode_length - 1) & (self.reset_buf != 0)
        self.timeout_buf = (self.progress_buf >= self.max_episode_length - 1) | (
            self.reset_buf != 0
        )

        # randomize observations
        if self.dr_randomizations.get("observations", None):
            self.obs_buf = self.dr_randomizations["observations"]["noise_lambda"](
                self.obs_buf
            )

        self.extras["time_outs"] = self.timeout_buf.to(self.rl_device)
        self.extras["success"] = self.success_buf.sum().to(self.rl_device)

        self.obs_dict["obs"] = torch.clamp(
            self.obs_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_left"] = torch.clamp(
            self.digit_left_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_right"] = torch.clamp(
            self.digit_right_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["gt_extrinsic_contact"] = torch.clamp(
            self.gt_extrinsic_contact, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["pointclouds_t"] = torch.clamp(
            self.pointclouds_t, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return (
            self.obs_dict,
            self.rew_buf.to(self.rl_device),
            self.reset_buf.to(self.rl_device),
            self.extras,
        )

    def zero_actions(self) -> torch.Tensor:
        """Returns a buffer with zero actions.

        Returns:
            A buffer of zero torch actions
        """
        actions = torch.zeros(
            [self.num_envs, self.num_actions],
            dtype=torch.float32,
            device=self.rl_device,
        )

        return actions

    def reset_idx(self, env_idx):
        """Reset environment with indces in env_idx.
        Should be implemented in an environment class inherited from VecTask.
        """
        pass

    def reset(self):
        """Is called only once when environment starts to provide the first observations.
        Doesn't calculate observations. Actual reset and observation calculation need to be implemented by user.
        Returns:
            Observation dictionary
        """
        self.obs_dict["obs"] = torch.clamp(
            self.obs_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_left"] = torch.clamp(
            self.digit_left_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_right"] = torch.clamp(
            self.digit_right_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["gt_extrinsic_contact"] = torch.clamp(
            self.gt_extrinsic_contact, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["pointclouds_t"] = torch.clamp(
            self.pointclouds_t, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return self.obs_dict

    def reset_done(self):
        """Reset the environment.
        Returns:
            Observation dictionary, indices of environments being reset
        """
        done_env_ids = self.reset_buf.nonzero(as_tuple=False).flatten()
        if len(done_env_ids) > 0:
            self.reset_idx(done_env_ids)

        self.obs_dict["obs"] = torch.clamp(
            self.obs_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_left"] = torch.clamp(
            self.digit_left_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["digits_right"] = torch.clamp(
            self.digit_right_buf, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["gt_extrinsic_contact"] = torch.clamp(
            self.gt_extrinsic_contact, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)
        self.obs_dict["pointclouds_t"] = torch.clamp(
            self.pointclouds_t, -self.clip_obs, self.clip_obs
        ).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return self.obs_dict, done_env_ids

    def render(self, mode="rgb_array"):
        """Draw the frame to the viewer, and check for keyboard events."""
        if self.viewer:
            # check for window closed
            if self.gym.query_viewer_has_closed(self.viewer):
                sys.exit()

            # check for keyboard events
            for evt in self.gym.query_viewer_action_events(self.viewer):
                if evt.action == "QUIT" and evt.value > 0:
                    sys.exit()
                elif evt.action == "toggle_viewer_sync" and evt.value > 0:
                    self.enable_viewer_sync = not self.enable_viewer_sync
                elif evt.action == "record_frames" and evt.value > 0:
                    self.record_frames = not self.record_frames

            # fetch results
            if self.device != "cpu":
                self.gym.fetch_results(self.sim, True)

            # step graphics
            if self.enable_viewer_sync:
                self.gym.step_graphics(self.sim)
                self.gym.draw_viewer(self.viewer, self.sim, True)

                # Wait for dt to elapse in real time.
                # This synchronizes the physics simulation with the rendering rate.
                self.gym.sync_frame_time(self.sim)

                # it seems like in some cases sync_frame_time still results in higher-than-realtime framerate
                # this code will slow down the rendering to real time
                now = time.time()
                delta = now - self.last_frame_time
                if self.render_fps < 0:
                    # render at control frequency
                    render_dt = (
                        self.dt * self.control_freq_inv
                    )  # render every control step
                else:
                    render_dt = 1.0 / self.render_fps

                if delta < render_dt:
                    time.sleep(render_dt - delta)

                self.last_frame_time = time.time()

            else:
                self.gym.poll_viewer_events(self.viewer)

            if self.record_frames:
                if not os.path.isdir(self.record_frames_dir):
                    os.makedirs(self.record_frames_dir, exist_ok=True)

                self.gym.write_viewer_image_to_file(
                    self.viewer,
                    join(self.record_frames_dir, f"frame_{self.control_steps}.png"),
                )

            if self.virtual_display and mode == "rgb_array":
                img = self.virtual_display.grab()
                return np.array(img)

    def __parse_sim_params(
        self, physics_engine: str, config_sim: Dict[str, Any]
    ) -> gymapi.SimParams:
        """Parse the config dictionary for physics stepping settings.

        Args:
            physics_engine: which physics engine to use. "physx" or "flex"
            config_sim: dict of sim configuration parameters
        Returns
            IsaacGym SimParams object with updated settings.
        """
        sim_params = gymapi.SimParams()

        # check correct up-axis
        if config_sim["up_axis"] not in ["z", "y"]:
            msg = f"Invalid physics up-axis: {config_sim['up_axis']}"
            print(msg)
            raise ValueError(msg)

        # assign general sim parameters
        sim_params.dt = config_sim["dt"]
        sim_params.num_client_threads = config_sim.get("num_client_threads", 0)
        sim_params.use_gpu_pipeline = config_sim["use_gpu_pipeline"]
        sim_params.substeps = config_sim.get("substeps", 2)

        # assign up-axis
        if config_sim["up_axis"] == "z":
            sim_params.up_axis = gymapi.UP_AXIS_Z
        else:
            sim_params.up_axis = gymapi.UP_AXIS_Y

        # assign gravity
        sim_params.gravity = gymapi.Vec3(*config_sim["gravity"])

        # configure physics parameters
        if physics_engine == "physx":
            # set the parameters
            if "physx" in config_sim:
                for opt in config_sim["physx"].keys():
                    if opt == "contact_collection":
                        setattr(
                            sim_params.physx,
                            opt,
                            gymapi.ContactCollection(config_sim["physx"][opt]),
                        )
                    else:
                        setattr(sim_params.physx, opt, config_sim["physx"][opt])
        else:
            # set the parameters
            if "flex" in config_sim:
                for opt in config_sim["flex"].keys():
                    setattr(sim_params.flex, opt, config_sim["flex"][opt])

        # return the configured params
        return sim_params

    """
    Domain Randomization methods
    """

    def get_actor_params_info(self, dr_params: Dict[str, Any], env):
        """Generate a flat array of actor params, their names and ranges.

        Returns:
            The array
        """

        if "actor_params" not in dr_params:
            return None
        params = []
        names = []
        lows = []
        highs = []
        param_getters_map = get_property_getter_map(self.gym)
        for actor, actor_properties in dr_params["actor_params"].items():
            handle = self.gym.find_actor_handle(env, actor)
            for prop_name, prop_attrs in actor_properties.items():
                if prop_name == "color":
                    continue  # this is set randomly
                props = param_getters_map[prop_name](env, handle)
                if not isinstance(props, list):
                    props = [props]
                for prop_idx, prop in enumerate(props):
                    for attr, attr_randomization_params in prop_attrs.items():
                        name = prop_name + "_" + str(prop_idx) + "_" + attr
                        lo_hi = attr_randomization_params["range"]
                        distr = attr_randomization_params["distribution"]
                        if "uniform" not in distr:
                            lo_hi = (-1.0 * float("Inf"), float("Inf"))
                        if isinstance(prop, np.ndarray):
                            for attr_idx in range(prop[attr].shape[0]):
                                params.append(prop[attr][attr_idx])
                                names.append(name + "_" + str(attr_idx))
                                lows.append(lo_hi[0])
                                highs.append(lo_hi[1])
                        else:
                            params.append(getattr(prop, attr))
                            names.append(name)
                            lows.append(lo_hi[0])
                            highs.append(lo_hi[1])
        return params, names, lows, highs

    def apply_randomizations(self, dr_params):
        """Apply domain randomizations to the environment.

        Note that currently we can only apply randomizations only on resets, due to current PhysX limitations

        Args:
            dr_params: parameters for domain randomization to use.
        """

        # If we don't have a randomization frequency, randomize every step
        rand_freq = dr_params.get("frequency", 1)

        # First, determine what to randomize:
        #   - non-environment parameters when > frequency steps have passed since the last non-environment
        #   - physical environments in the reset buffer, which have exceeded the randomization frequency threshold
        #   - on the first call, randomize everything
        self.last_step = self.gym.get_frame_count(self.sim)
        if self.first_randomization:
            do_nonenv_randomize = True
            env_ids = list(range(self.num_envs))
        else:
            do_nonenv_randomize = (self.last_step - self.last_rand_step) >= rand_freq
            rand_envs = torch.where(
                self.randomize_buf >= rand_freq,
                torch.ones_like(self.randomize_buf),
                torch.zeros_like(self.randomize_buf),
            )
            rand_envs = torch.logical_and(rand_envs, self.reset_buf)
            env_ids = torch.nonzero(rand_envs, as_tuple=False).squeeze(-1).tolist()
            self.randomize_buf[rand_envs] = 0

        if do_nonenv_randomize:
            self.last_rand_step = self.last_step

        param_setters_map = get_property_setter_map(self.gym)
        param_setter_defaults_map = get_default_setter_args(self.gym)
        param_getters_map = get_property_getter_map(self.gym)

        # On first iteration, check the number of buckets
        if self.first_randomization:
            check_buckets(self.gym, self.envs, dr_params)

        for nonphysical_param in ["observations", "actions"]:
            if nonphysical_param in dr_params and do_nonenv_randomize:
                dist = dr_params[nonphysical_param]["distribution"]
                op_type = dr_params[nonphysical_param]["operation"]
                sched_type = (
                    dr_params[nonphysical_param]["schedule"]
                    if "schedule" in dr_params[nonphysical_param]
                    else None
                )
                sched_step = (
                    dr_params[nonphysical_param]["schedule_steps"]
                    if "schedule" in dr_params[nonphysical_param]
                    else None
                )
                op = operator.add if op_type == "additive" else operator.mul

                if sched_type == "linear":
                    sched_scaling = 1.0 / sched_step * min(self.last_step, sched_step)
                elif sched_type == "constant":
                    sched_scaling = 0 if self.last_step < sched_step else 1
                else:
                    sched_scaling = 1

                if dist == "gaussian":
                    mu, var = dr_params[nonphysical_param]["range"]
                    mu_corr, var_corr = dr_params[nonphysical_param].get(
                        "range_correlated", [0.0, 0.0]
                    )

                    if op_type == "additive":
                        mu *= sched_scaling
                        var *= sched_scaling
                        mu_corr *= sched_scaling
                        var_corr *= sched_scaling
                    elif op_type == "scaling":
                        var = var * sched_scaling  # scale up var over time
                        mu = mu * sched_scaling + 1.0 * (
                            1.0 - sched_scaling
                        )  # linearly interpolate

                        var_corr = var_corr * sched_scaling  # scale up var over time
                        mu_corr = mu_corr * sched_scaling + 1.0 * (
                            1.0 - sched_scaling
                        )  # linearly interpolate

                    def noise_lambda(tensor, param_name=nonphysical_param):
                        params = self.dr_randomizations[param_name]
                        corr = params.get("corr", None)
                        if corr is None:
                            corr = torch.randn_like(tensor)
                            params["corr"] = corr
                        corr = corr * params["var_corr"] + params["mu_corr"]
                        return op(
                            tensor,
                            corr
                            + torch.randn_like(tensor) * params["var"]
                            + params["mu"],
                        )

                    self.dr_randomizations[nonphysical_param] = {
                        "mu": mu,
                        "var": var,
                        "mu_corr": mu_corr,
                        "var_corr": var_corr,
                        "noise_lambda": noise_lambda,
                    }

                elif dist == "uniform":
                    lo, hi = dr_params[nonphysical_param]["range"]
                    lo_corr, hi_corr = dr_params[nonphysical_param].get(
                        "range_correlated", [0.0, 0.0]
                    )

                    if op_type == "additive":
                        lo *= sched_scaling
                        hi *= sched_scaling
                        lo_corr *= sched_scaling
                        hi_corr *= sched_scaling
                    elif op_type == "scaling":
                        lo = lo * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        hi = hi * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        lo_corr = lo_corr * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        hi_corr = hi_corr * sched_scaling + 1.0 * (1.0 - sched_scaling)

                    def noise_lambda(tensor, param_name=nonphysical_param):
                        params = self.dr_randomizations[param_name]
                        corr = params.get("corr", None)
                        if corr is None:
                            corr = torch.randn_like(tensor)
                            params["corr"] = corr
                        corr = (
                            corr * (params["hi_corr"] - params["lo_corr"])
                            + params["lo_corr"]
                        )
                        return op(
                            tensor,
                            corr
                            + torch.rand_like(tensor) * (params["hi"] - params["lo"])
                            + params["lo"],
                        )

                    self.dr_randomizations[nonphysical_param] = {
                        "lo": lo,
                        "hi": hi,
                        "lo_corr": lo_corr,
                        "hi_corr": hi_corr,
                        "noise_lambda": noise_lambda,
                    }

        if "sim_params" in dr_params and do_nonenv_randomize:
            prop_attrs = dr_params["sim_params"]
            prop = self.gym.get_sim_params(self.sim)

            if self.first_randomization:
                self.original_props["sim_params"] = {
                    attr: getattr(prop, attr) for attr in dir(prop)
                }

            for attr, attr_randomization_params in prop_attrs.items():
                apply_random_samples(
                    prop,
                    self.original_props["sim_params"],
                    attr,
                    attr_randomization_params,
                    self.last_step,
                )

            self.gym.set_sim_params(self.sim, prop)

        # If self.actor_params_generator is initialized: use it to
        # sample actor simulation params. This gives users the
        # freedom to generate samples from arbitrary distributions,
        # e.g. use full-covariance distributions instead of the DR's
        # default of treating each simulation parameter independently.
        extern_offsets = {}
        if self.actor_params_generator is not None:
            for env_id in env_ids:
                self.extern_actor_params[env_id] = self.actor_params_generator.sample()
                extern_offsets[env_id] = 0

        # randomise all attributes of each actor (hand, cube etc..)
        # actor_properties are (stiffness, damping etc..)

        # Loop over actors, then loop over envs, then loop over their props
        # and lastly loop over the ranges of the params

        for actor, actor_properties in dr_params["actor_params"].items():
            # Loop over all envs as this part is not tensorised yet
            for env_id in env_ids:
                env = self.envs[env_id]
                handle = self.gym.find_actor_handle(env, actor)
                extern_sample = self.extern_actor_params[env_id]

                # randomise dof_props, rigid_body, rigid_shape properties
                # all obtained from the YAML file
                # EXAMPLE: prop name: dof_properties, rigid_body_properties, rigid_shape properties
                #          prop_attrs:
                #               {'damping': {'range': [0.3, 3.0], 'operation': 'scaling', 'distribution': 'loguniform'}
                #               {'stiffness': {'range': [0.75, 1.5], 'operation': 'scaling', 'distribution': 'loguniform'}
                for prop_name, prop_attrs in actor_properties.items():
                    if prop_name == "color":
                        num_bodies = self.gym.get_actor_rigid_body_count(env, handle)
                        for n in range(num_bodies):
                            self.gym.set_rigid_body_color(
                                env,
                                handle,
                                n,
                                gymapi.MESH_VISUAL,
                                gymapi.Vec3(
                                    random.uniform(0, 1),
                                    random.uniform(0, 1),
                                    random.uniform(0, 1),
                                ),
                            )
                        continue

                    if prop_name == "scale":
                        setup_only = prop_attrs.get("setup_only", False)
                        if (setup_only and not self.sim_initialized) or not setup_only:
                            attr_randomization_params = prop_attrs
                            sample = generate_random_samples(
                                attr_randomization_params, 1, self.last_step, None
                            )
                            og_scale = 1
                            if attr_randomization_params["operation"] == "scaling":
                                new_scale = og_scale * sample
                            elif attr_randomization_params["operation"] == "additive":
                                new_scale = og_scale + sample
                            self.gym.set_actor_scale(env, handle, new_scale)
                        continue

                    prop = param_getters_map[prop_name](env, handle)
                    set_random_properties = True

                    if isinstance(prop, list):
                        if self.first_randomization:
                            self.original_props[prop_name] = [
                                {attr: getattr(p, attr) for attr in dir(p)}
                                for p in prop
                            ]
                        for p, og_p in zip(prop, self.original_props[prop_name]):
                            for attr, attr_randomization_params in prop_attrs.items():
                                setup_only = attr_randomization_params.get(
                                    "setup_only", False
                                )
                                if (
                                    setup_only and not self.sim_initialized
                                ) or not setup_only:
                                    smpl = None
                                    if self.actor_params_generator is not None:
                                        (
                                            smpl,
                                            extern_offsets[env_id],
                                        ) = get_attr_val_from_sample(
                                            extern_sample,
                                            extern_offsets[env_id],
                                            p,
                                            attr,
                                        )
                                    apply_random_samples(
                                        p,
                                        og_p,
                                        attr,
                                        attr_randomization_params,
                                        self.last_step,
                                        smpl,
                                    )
                                else:
                                    set_random_properties = False
                    else:
                        if self.first_randomization:
                            self.original_props[prop_name] = deepcopy(prop)
                        for attr, attr_randomization_params in prop_attrs.items():
                            setup_only = attr_randomization_params.get(
                                "setup_only", False
                            )
                            if (
                                setup_only and not self.sim_initialized
                            ) or not setup_only:
                                smpl = None
                                if self.actor_params_generator is not None:
                                    (
                                        smpl,
                                        extern_offsets[env_id],
                                    ) = get_attr_val_from_sample(
                                        extern_sample,
                                        extern_offsets[env_id],
                                        prop,
                                        attr,
                                    )
                                apply_random_samples(
                                    prop,
                                    self.original_props[prop_name],
                                    attr,
                                    attr_randomization_params,
                                    self.last_step,
                                    smpl,
                                )
                            else:
                                set_random_properties = False

                    if set_random_properties:
                        setter = param_setters_map[prop_name]
                        default_args = param_setter_defaults_map[prop_name]
                        setter(env, handle, prop, *default_args)

        if self.actor_params_generator is not None:
            for env_id in env_ids:  # check that we used all dims in sample
                if extern_offsets[env_id] > 0:
                    extern_sample = self.extern_actor_params[env_id]
                    if extern_offsets[env_id] != extern_sample.shape[0]:
                        print(
                            "env_id",
                            env_id,
                            "extern_offset",
                            extern_offsets[env_id],
                            "vs extern_sample.shape",
                            extern_sample.shape,
                        )
                        raise Exception("Invalid extern_sample size")

        self.first_randomization = False

NCF_policies/isaacgymenvs/tasks/base/vec_task.py
carolinahiguera/NCF-Policy
# --------------------------------------------------------
# Monte-Carlo inference over distributions across sliding touch 
# https://arxiv.org/abs/2210.14210
# https://github.com/facebookresearch/MidasTouch
# Sudharshan Suresh, Zilin Si, Stuart Anderson, Michael Kaess, Mustafa Mukadam
# --------------------------------------------------------

from os import path as osp
import numpy as np
import tacto
import cv2


import torch
from scipy.spatial.transform import Rotation as R
import trimesh
import random

DEBUG = False


class digit_renderer:
    def __init__(
        self,
        cfg,
        obj_path: str = None,
        obj_quat: np.ndarray = [0.0, 0.0, 0.0, 1.0],
        obj_scale: float = 1.0,
        headless=False,
        bg_id=None,
    ):
        self.render_config = cfg
        self.randomize = self.render_config.randomize
        if self.randomize:
            bg_id = random.randint(0, 9)
        elif bg_id is None:
            bg_id = cfg.bg_id

        # Create renderer
        self.renderer = tacto.Renderer(
            width=cfg.width,
            height=cfg.height,
            background=cv2.imread(tacto.get_background_image_path(bg_id)),
            config_path=tacto.get_digit_shadow_config_path(),
            # config_path=tacto.get_digit_config_path(),
            headless=headless,
        )
        self.cam_dist = cfg.cam_dist
        self.pixmm = cfg.pixmm

        if not DEBUG:
            _, self.bg_depth = self.renderer.render()
            self.bg_depth = self.bg_depth[0]
            self.bg_depth_pix = self.correct_pyrender_height_map(self.bg_depth)

        if obj_path is not None:
            # self.obj_loader = object_loader(obj_path)
            # self.obj_loader.set_object_pose(pose=obj_pose)
            obj_trimesh = trimesh.load(obj_path)
            obj_trimesh = obj_trimesh.apply_scale(obj_scale)
            self.obj_mesh = obj_trimesh
            obj_euler = R.from_quat(obj_quat).as_euler("xyz", degrees=False)
            self.renderer.add_object(obj_trimesh, "object", orientation=obj_euler)

        self.press_depth = 0.001

    def randomize_light(self):
        self.renderer.randomize_light()

    def get_background(self, frame="gel"):
        """
        Return cached bg image
        """
        return self.bg_depth_pix if frame == "gel" else self.bg_depth

    def pix2meter(self, pix):
        """
        Convert pixel to meter
        """
        return pix * self.pixmm / 1000.0

    def meter2pix(self, m):
        """
        Convert meter to pixels
        """
        return m * 1000.0 / self.pixmm

    def update_object_pose(self, pose):
        # (x, y, z) and euler angles
        self.renderer.update_object_pose(obj_name="object", pose=pose)

    def update_pose_given_point(self, point, press_depth, shear_mag, delta):
        """
        Convert meter to pixels
        """
        dist = np.linalg.norm(point - self.obj_loader.obj_vertices, axis=1)
        idx = np.argmin(dist)

        # idx: the idx vertice, get a new pose
        new_position = self.obj_loader.obj_vertices[idx].copy()
        new_orientation = self.obj_loader.obj_normals[idx].copy()

        delta = np.random.uniform(low=0.0, high=2 * np.pi, size=(1,))[0]
        new_pose = pose_from_vertex_normal(
            new_position, new_orientation, shear_mag, delta
        ).squeeze()
        self.update_pose_given_pose(press_depth, new_pose)

    def update_pose_given_pose(self, cam_pose):
        """
        Given tf gel_pose and press_depth, update tacto camera
        """
        self.renderer.update_camera_pose_from_matrix(cam_pose)

    def add_press(self, pose):
        """
        Add sensor penetration
        """
        pen_mat = np.eye(4)
        pen_mat[0, 3] = self.press_depth
        return np.matmul(pose, pen_mat)

    def gel2cam(self, gel_pose):
        """
        Convert gel_pose to cam_pose
        """
        cam_tf = np.eye(4)
        cam_tf[2, 3] = self.cam_dist
        return np.matmul(gel_pose, cam_tf)

    def cam2gel(self, cam_pose):
        """
        Convert cam_pose to gel_pose
        """
        cam_tf = np.eye(4)
        cam_tf[2, 3] = -self.cam_dist
        return np.matmul(cam_pose, cam_tf)

    # input depth is in camera frame here
    def render(self):
        """
        render [tactile image + depth + mask] @ current pose
        """
        color, depth = self.renderer.render()
        color, depth = color[0], depth[0]
        diff_depth = (self.bg_depth) - depth
        contact_mask = diff_depth > np.abs(self.press_depth * 0.2)
        gel_depth = self.correct_pyrender_height_map(depth)  #  pix in gel frame
        # cam_depth = self.correct_image_height_map(gel_depth) #  pix in gel frame
        # assert np.allclose(cam_depth, depth), "Conversion to pixels is incorrect"
        if self.randomize:
            self.renderer.randomize_light()
        return color, gel_depth, contact_mask

    def correct_pyrender_height_map(self, height_map):
        """
        Input: height_map in meters, in camera frame
        Output: height_map in pixels, in gel frame
        """
        # move to the gel center
        height_map = (self.cam_dist - height_map) * (1000 / self.pixmm)
        return height_map

    def correct_image_height_map(self, height_map, output_frame="cam"):
        """
        Input: height_map in pixels, in gel frame
        Output: height_map in meters, in camera/gel frame
        """
        height_map = (
            -height_map * (self.pixmm / 1000)
            + float(output_frame == "cam") * self.cam_dist
        )
        return height_map

    def get_cam_pose_matrix(self):
        """
        return camera pose matrix of renderer
        """
        return self.renderer.camera_nodes[0].matrix

    def get_cam_pose(self):
        """
        return camera pose of renderer
        """
        # print(f"Cam pose: {tf_to_xyzquat(self.get_cam_pose_matrix())}")
        return self.get_cam_pose_matrix()

    def get_gel_pose_matrix(self):
        """
        return gel pose matrix of renderer
        """
        return self.cam2gel(self.get_cam_pose_matrix())

    def get_gel_pose(self):
        """
        return gel pose of renderer
        """
        # print(f"Gel pose: {tf_to_xyzquat(self.get_gel_pose_matrix())}")
        return self.get_gel_pose_matrix()

    def heightmap2Pointcloud(self, depth, contact_mask=None):
        """
        Convert heightmap + contact mask to point cloud
        [Input]  depth: (width, height) in pixels, in gel frame, Contact mask: binary (width, height)
        [Output] pointcloud: [(width, height) - (masked off points), 3] in meters in camera frame
        """
        depth = self.correct_image_height_map(depth, output_frame="cam")

        if contact_mask is not None:
            heightmapValid = depth * contact_mask  # apply contact mask
        else:
            heightmapValid = depth

        f, w, h = self.renderer.f, self.renderer.width / 2.0, self.renderer.height / 2.0

        if not torch.is_tensor(heightmapValid):
            heightmapValid = torch.from_numpy(heightmapValid)
        # (0, 640) and (0, 480)
        xvals = torch.arange(heightmapValid.shape[1], device=heightmapValid.device)
        yvals = torch.arange(heightmapValid.shape[0], device=heightmapValid.device)
        [x, y] = torch.meshgrid(xvals, yvals)
        x, y = torch.transpose(x, 0, 1), torch.transpose(
            y, 0, 1
        )  # future warning: https://github.com/pytorch/pytorch/issues/50276

        # x and y in meters
        x = ((x - w)) / f
        y = ((y - h)) / f

        x *= depth
        y *= -depth

        heightmap_3d = torch.hstack(
            (x.reshape((-1, 1)), y.reshape((-1, 1)), heightmapValid.reshape((-1, 1)))
        )

        heightmap_3d[:, 2] *= -1
        heightmap_3d = heightmap_3d[heightmap_3d[:, 2] != 0]
        return heightmap_3d
NCF_policies/isaacgymenvs/tasks/ncf_tacto/digit_renderer.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: base class.

Inherits Gym's VecTask class and abstract base class. Inherited by environment classes. Not directly executed.

Configuration defined in FactoryBase.yaml. Asset info defined in factory_asset_info_franka_table.yaml.
"""


import hydra
import math
import numpy as np
import os
import sys
import torch

from gym import logger
from isaacgym import gymapi, gymtorch
from isaacgymenvs.utils import torch_jit_utils as torch_utils
from isaacgymenvs.tasks.base.vec_task import VecTask
import isaacgymenvs.tasks.ncf_tacto.factory_control as fc
from isaacgymenvs.tasks.ncf_tacto.factory_schema_class_base import FactoryABCBase
from isaacgymenvs.tasks.ncf_tacto.factory_schema_config_base import (
    FactorySchemaConfigBase,
)


class FactoryBase(VecTask, FactoryABCBase):
    def __init__(
        self,
        cfg,
        rl_device,
        sim_device,
        graphics_device_id,
        headless,
        virtual_screen_capture,
        force_render,
    ):
        """Initialize instance variables. Initialize VecTask superclass."""

        self.cfg = cfg
        self.cfg["headless"] = headless

        self._get_base_yaml_params()

        if self.cfg_base.mode.export_scene:
            sim_device = "cpu"

        super().__init__(
            cfg,
            rl_device,
            sim_device,
            graphics_device_id,
            headless,
            virtual_screen_capture,
            force_render,
        )  # create_sim() is called here

    def _get_base_yaml_params(self):
        """Initialize instance variables from YAML files."""

        cs = hydra.core.config_store.ConfigStore.instance()
        cs.store(name="factory_schema_config_base", node=FactorySchemaConfigBase)

        config_path = (
            "task/FactoryBase.yaml"  # relative to Gym's Hydra search path (cfg dir)
        )
        self.cfg_base = hydra.compose(config_name=config_path)
        self.cfg_base = self.cfg_base["task"]  # strip superfluous nesting

        asset_info_path = "../../assets/factory/yaml/factory_asset_info_franka_table.yaml"  # relative to Gym's Hydra search path (cfg dir)
        self.asset_info_franka_table = hydra.compose(config_name=asset_info_path)
        self.asset_info_franka_table = self.asset_info_franka_table[""][""][""][""][""][
            ""
        ]["assets"]["factory"][
            "yaml"
        ]  # strip superfluous nesting

    def create_sim(self):
        """Set sim and PhysX params. Create sim object, ground plane, and envs."""

        if self.cfg_base.mode.export_scene:
            self.sim_params.use_gpu_pipeline = False

        self.sim = super().create_sim(
            compute_device=self.device_id,
            graphics_device=self.graphics_device_id,
            physics_engine=self.physics_engine,
            sim_params=self.sim_params,
        )
        self._create_ground_plane()
        self.create_envs()  # defined in subclass

    def _create_ground_plane(self):
        """Set ground plane params. Add plane."""

        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0)
        plane_params.distance = 0.0  # default = 0.0
        plane_params.static_friction = 1.0  # default = 1.0
        plane_params.dynamic_friction = 1.0  # default = 1.0
        plane_params.restitution = 0.0  # default = 0.0

        # self.gym.add_ground(self.sim, plane_params)

    def import_franka_assets(self):
        """Set Franka and table asset options. Import assets."""

        urdf_root = os.path.join(
            os.path.dirname(__file__), "..", "..", "..", "assets", "factory", "urdf"
        )
        franka_file = "factory_franka_digit.urdf"

        franka_options = gymapi.AssetOptions()
        franka_options.flip_visual_attachments = True
        franka_options.fix_base_link = True
        franka_options.collapse_fixed_joints = False
        franka_options.thickness = 0.0  # default = 0.02
        franka_options.density = 1000.0  # default = 1000.0
        franka_options.armature = 0.01  # default = 0.0
        franka_options.use_physx_armature = True
        if self.cfg_base.sim.add_damping:
            franka_options.linear_damping = (
                1.0  # default = 0.0; increased to improve stability
            )
            franka_options.max_linear_velocity = (
                1.0  # default = 1000.0; reduced to prevent CUDA errors
            )
            franka_options.angular_damping = (
                5.0  # default = 0.5; increased to improve stability
            )
            franka_options.max_angular_velocity = (
                2 * math.pi
            )  # default = 64.0; reduced to prevent CUDA errors
        else:
            franka_options.linear_damping = 0.0  # default = 0.0
            franka_options.max_linear_velocity = 1000.0  # default = 1000.0
            franka_options.angular_damping = 0.5  # default = 0.5
            franka_options.max_angular_velocity = 64.0  # default = 64.0
        franka_options.disable_gravity = True
        franka_options.enable_gyroscopic_forces = True
        franka_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        franka_options.use_mesh_materials = True
        if self.cfg_base.mode.export_scene:
            franka_options.mesh_normal_mode = gymapi.COMPUTE_PER_FACE

        table_options = gymapi.AssetOptions()
        table_options.flip_visual_attachments = False  # default = False
        table_options.fix_base_link = True
        table_options.thickness = 0.0  # default = 0.02
        table_options.density = 1000.0  # default = 1000.0
        table_options.armature = 0.0  # default = 0.0
        table_options.use_physx_armature = True
        table_options.linear_damping = 0.0  # default = 0.0
        table_options.max_linear_velocity = 1000.0  # default = 1000.0
        table_options.angular_damping = 0.0  # default = 0.5
        table_options.max_angular_velocity = 64.0  # default = 64.0
        table_options.disable_gravity = False
        table_options.enable_gyroscopic_forces = True
        table_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        table_options.use_mesh_materials = False
        if self.cfg_base.mode.export_scene:
            table_options.mesh_normal_mode = gymapi.COMPUTE_PER_FACE

        franka_asset = self.gym.load_asset(
            self.sim, urdf_root, franka_file, franka_options
        )
        table_asset = self.gym.create_box(
            self.sim,
            self.asset_info_franka_table.table_depth,
            self.asset_info_franka_table.table_width,
            self.cfg_base.env.table_height,
            table_options,
        )

        return franka_asset, table_asset

    def acquire_base_tensors(self):
        """Acquire and wrap tensors. Create views."""

        _root_state = self.gym.acquire_actor_root_state_tensor(
            self.sim
        )  # shape = (num_envs * num_actors, 13)
        _body_state = self.gym.acquire_rigid_body_state_tensor(
            self.sim
        )  # shape = (num_envs * num_bodies, 13)
        _dof_state = self.gym.acquire_dof_state_tensor(
            self.sim
        )  # shape = (num_envs * num_dofs, 2)
        _dof_force = self.gym.acquire_dof_force_tensor(
            self.sim
        )  # shape = (num_envs * num_dofs, 1)
        _contact_force = self.gym.acquire_net_contact_force_tensor(
            self.sim
        )  # shape = (num_envs * num_bodies, 3)
        _jacobian = self.gym.acquire_jacobian_tensor(
            self.sim, "franka"
        )  # shape = (num envs, num_bodies, 6, num_dofs)
        _mass_matrix = self.gym.acquire_mass_matrix_tensor(
            self.sim, "franka"
        )  # shape = (num_envs, num_dofs, num_dofs)

        self.root_state = gymtorch.wrap_tensor(_root_state)
        self.body_state = gymtorch.wrap_tensor(_body_state)
        self.dof_state = gymtorch.wrap_tensor(_dof_state)
        self.dof_force = gymtorch.wrap_tensor(_dof_force)
        self.contact_force = gymtorch.wrap_tensor(_contact_force)
        self.jacobian = gymtorch.wrap_tensor(_jacobian)
        self.mass_matrix = gymtorch.wrap_tensor(_mass_matrix)

        self.root_pos = self.root_state.view(self.num_envs, self.num_actors, 13)[
            ..., 0:3
        ]
        self.root_quat = self.root_state.view(self.num_envs, self.num_actors, 13)[
            ..., 3:7
        ]
        self.root_linvel = self.root_state.view(self.num_envs, self.num_actors, 13)[
            ..., 7:10
        ]
        self.root_angvel = self.root_state.view(self.num_envs, self.num_actors, 13)[
            ..., 10:13
        ]
        self.body_pos = self.body_state.view(self.num_envs, self.num_bodies, 13)[
            ..., 0:3
        ]
        self.body_quat = self.body_state.view(self.num_envs, self.num_bodies, 13)[
            ..., 3:7
        ]
        self.body_linvel = self.body_state.view(self.num_envs, self.num_bodies, 13)[
            ..., 7:10
        ]
        self.body_angvel = self.body_state.view(self.num_envs, self.num_bodies, 13)[
            ..., 10:13
        ]
        self.dof_pos = self.dof_state.view(self.num_envs, self.num_dofs, 2)[..., 0]
        self.dof_vel = self.dof_state.view(self.num_envs, self.num_dofs, 2)[..., 1]
        self.dof_force_view = self.dof_force.view(self.num_envs, self.num_dofs, 1)[
            ..., 0
        ]
        self.contact_force = self.contact_force.view(self.num_envs, self.num_bodies, 3)[
            ..., 0:3
        ]

        self.arm_dof_pos = self.dof_pos[:, 0:7]
        self.arm_mass_matrix = self.mass_matrix[
            :, 0:7, 0:7
        ]  # for Franka arm (not gripper)

        self.hand_pos = self.body_pos[:, self.hand_body_id_env, 0:3]
        self.hand_quat = self.body_quat[:, self.hand_body_id_env, 0:4]
        self.hand_linvel = self.body_linvel[:, self.hand_body_id_env, 0:3]
        self.hand_angvel = self.body_angvel[:, self.hand_body_id_env, 0:3]
        self.hand_jacobian = self.jacobian[
            :, self.hand_body_id_env - 1, 0:6, 0:7
        ]  # minus 1 because base is fixed

        self.left_finger_pos = self.body_pos[:, self.left_finger_body_id_env, 0:3]
        self.left_finger_quat = self.body_quat[:, self.left_finger_body_id_env, 0:4]
        self.left_finger_linvel = self.body_linvel[:, self.left_finger_body_id_env, 0:3]
        self.left_finger_angvel = self.body_angvel[:, self.left_finger_body_id_env, 0:3]
        self.left_finger_jacobian = self.jacobian[
            :, self.left_finger_body_id_env - 1, 0:6, 0:7
        ]  # minus 1 because base is fixed

        self.right_finger_pos = self.body_pos[:, self.right_finger_body_id_env, 0:3]
        self.right_finger_quat = self.body_quat[:, self.right_finger_body_id_env, 0:4]
        self.right_finger_linvel = self.body_linvel[
            :, self.right_finger_body_id_env, 0:3
        ]
        self.right_finger_angvel = self.body_angvel[
            :, self.right_finger_body_id_env, 0:3
        ]
        self.right_finger_jacobian = self.jacobian[
            :, self.right_finger_body_id_env - 1, 0:6, 0:7
        ]  # minus 1 because base is fixed

        self.left_finger_force = self.contact_force[
            :, self.left_finger_body_id_env, 0:3
        ]
        self.right_finger_force = self.contact_force[
            :, self.right_finger_body_id_env, 0:3
        ]

        self.left_fingertip_pos = self.body_pos[:, self.left_fingertip_body_id_env, 0:3]
        self.left_fingertip_quat = self.body_quat[
            :, self.left_fingertip_body_id_env, 0:4
        ]

        self.right_fingertip_pos = self.body_pos[
            :, self.right_fingertip_body_id_env, 0:3
        ]
        self.right_fingertip_quat = self.body_quat[
            :, self.right_fingertip_body_id_env, 0:4
        ]

        self.gripper_dof_pos = self.dof_pos[:, 7:9]

        self.fingertip_centered_pos = self.body_pos[
            :, self.fingertip_centered_body_id_env, 0:3
        ]
        self.fingertip_centered_quat = self.body_quat[
            :, self.fingertip_centered_body_id_env, 0:4
        ]
        self.fingertip_centered_linvel = self.body_linvel[
            :, self.fingertip_centered_body_id_env, 0:3
        ]
        self.fingertip_centered_angvel = self.body_angvel[
            :, self.fingertip_centered_body_id_env, 0:3
        ]
        self.fingertip_centered_jacobian = self.jacobian[
            :, self.fingertip_centered_body_id_env - 1, 0:6, 0:7
        ]  # minus 1 because base is fixed

        self.fingertip_midpoint_pos = (
            self.fingertip_centered_pos.detach().clone()
        )  # initial value
        self.fingertip_midpoint_quat = self.fingertip_centered_quat  # always equal
        self.fingertip_midpoint_linvel = (
            self.fingertip_centered_linvel.detach().clone()
        )  # initial value
        # From sum of angular velocities (https://physics.stackexchange.com/questions/547698/understanding-addition-of-angular-velocity),
        # angular velocity of midpoint w.r.t. world is equal to sum of
        # angular velocity of midpoint w.r.t. hand and angular velocity of hand w.r.t. world.
        # Midpoint is in sliding contact (i.e., linear relative motion) with hand; angular velocity of midpoint w.r.t. hand is zero.
        # Thus, angular velocity of midpoint w.r.t. world is equal to angular velocity of hand w.r.t. world.
        self.fingertip_midpoint_angvel = self.fingertip_centered_angvel  # always equal
        self.fingertip_midpoint_jacobian = (
            self.left_finger_jacobian + self.right_finger_jacobian
        ) * 0.5  # approximation

        self.dof_torque = torch.zeros(
            (self.num_envs, self.num_dofs), device=self.device
        )
        self.fingertip_contact_wrench = torch.zeros(
            (self.num_envs, 6), device=self.device
        )

        self.ctrl_target_fingertip_midpoint_pos = torch.zeros(
            (self.num_envs, 3), device=self.device
        )
        self.ctrl_target_fingertip_midpoint_quat = torch.zeros(
            (self.num_envs, 4), device=self.device
        )
        self.ctrl_target_dof_pos = torch.zeros(
            (self.num_envs, self.num_dofs), device=self.device
        )
        self.ctrl_target_gripper_dof_pos = torch.zeros(
            (self.num_envs, 2), device=self.device
        )
        self.ctrl_target_fingertip_contact_wrench = torch.zeros(
            (self.num_envs, 6), device=self.device
        )

        self.prev_actions = torch.zeros(
            (self.num_envs, self.num_actions), device=self.device
        )

    def refresh_base_tensors(self):
        """Refresh tensors."""
        # NOTE: Tensor refresh functions should be called once per step, before setters.

        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_rigid_body_state_tensor(self.sim)
        self.gym.refresh_dof_force_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)
        self.gym.refresh_jacobian_tensors(self.sim)
        self.gym.refresh_mass_matrix_tensors(self.sim)

        self.finger_midpoint_pos = (self.left_finger_pos + self.right_finger_pos) * 0.5
        self.fingertip_midpoint_pos = fc.translate_along_local_z(
            pos=self.finger_midpoint_pos,
            quat=self.hand_quat,
            offset=self.asset_info_franka_table.franka_finger_length,
            device=self.device,
        )
        # TODO: Add relative velocity term (see https://dynamicsmotioncontrol487379916.files.wordpress.com/2020/11/21-me258pointmovingrigidbody.pdf)
        self.fingertip_midpoint_linvel = self.fingertip_centered_linvel + torch.cross(
            self.fingertip_centered_angvel,
            (self.fingertip_midpoint_pos - self.fingertip_centered_pos),
            dim=1,
        )
        self.fingertip_midpoint_jacobian = (
            self.left_finger_jacobian + self.right_finger_jacobian
        ) * 0.5  # approximation

    def parse_controller_spec(self):
        """Parse controller specification into lower-level controller configuration."""

        cfg_ctrl_keys = {
            "num_envs",
            "jacobian_type",
            "gripper_prop_gains",
            "gripper_deriv_gains",
            "motor_ctrl_mode",
            "gain_space",
            "ik_method",
            "joint_prop_gains",
            "joint_deriv_gains",
            "do_motion_ctrl",
            "task_prop_gains",
            "task_deriv_gains",
            "do_inertial_comp",
            "motion_ctrl_axes",
            "do_force_ctrl",
            "force_ctrl_method",
            "wrench_prop_gains",
            "force_ctrl_axes",
        }
        self.cfg_ctrl = {cfg_ctrl_key: None for cfg_ctrl_key in cfg_ctrl_keys}

        self.cfg_ctrl["num_envs"] = self.num_envs
        self.cfg_ctrl["jacobian_type"] = self.cfg_task.ctrl.all.jacobian_type
        self.cfg_ctrl["gripper_prop_gains"] = torch.tensor(
            self.cfg_task.ctrl.all.gripper_prop_gains, device=self.device
        ).repeat((self.num_envs, 1))
        self.cfg_ctrl["gripper_deriv_gains"] = torch.tensor(
            self.cfg_task.ctrl.all.gripper_deriv_gains, device=self.device
        ).repeat((self.num_envs, 1))

        ctrl_type = self.cfg_task.ctrl.ctrl_type
        if ctrl_type == "gym_default":
            self.cfg_ctrl["motor_ctrl_mode"] = "gym"
            self.cfg_ctrl["gain_space"] = "joint"
            self.cfg_ctrl["ik_method"] = self.cfg_task.ctrl.gym_default.ik_method
            self.cfg_ctrl["joint_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.gym_default.joint_prop_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["joint_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.gym_default.joint_deriv_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["gripper_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.gym_default.gripper_prop_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["gripper_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.gym_default.gripper_deriv_gains, device=self.device
            ).repeat((self.num_envs, 1))
        elif ctrl_type == "joint_space_ik":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "joint"
            self.cfg_ctrl["ik_method"] = self.cfg_task.ctrl.joint_space_ik.ik_method
            self.cfg_ctrl["joint_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.joint_space_ik.joint_prop_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["joint_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.joint_space_ik.joint_deriv_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_inertial_comp"] = False
        elif ctrl_type == "joint_space_id":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "joint"
            self.cfg_ctrl["ik_method"] = self.cfg_task.ctrl.joint_space_id.ik_method
            self.cfg_ctrl["joint_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.joint_space_id.joint_prop_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["joint_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.joint_space_id.joint_deriv_gains, device=self.device
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_inertial_comp"] = True
        elif ctrl_type == "task_space_impedance":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "task"
            self.cfg_ctrl["do_motion_ctrl"] = True
            self.cfg_ctrl["task_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.task_space_impedance.task_prop_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["task_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.task_space_impedance.task_deriv_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_inertial_comp"] = False
            self.cfg_ctrl["motion_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.task_space_impedance.motion_ctrl_axes,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_force_ctrl"] = False
        elif ctrl_type == "operational_space_motion":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "task"
            self.cfg_ctrl["do_motion_ctrl"] = True
            self.cfg_ctrl["task_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.operational_space_motion.task_prop_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["task_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.operational_space_motion.task_deriv_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_inertial_comp"] = True
            self.cfg_ctrl["motion_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.operational_space_motion.motion_ctrl_axes,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_force_ctrl"] = False
        elif ctrl_type == "open_loop_force":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "task"
            self.cfg_ctrl["do_motion_ctrl"] = False
            self.cfg_ctrl["do_force_ctrl"] = True
            self.cfg_ctrl["force_ctrl_method"] = "open"
            self.cfg_ctrl["force_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.open_loop_force.force_ctrl_axes, device=self.device
            ).repeat((self.num_envs, 1))
        elif ctrl_type == "closed_loop_force":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "task"
            self.cfg_ctrl["do_motion_ctrl"] = False
            self.cfg_ctrl["do_force_ctrl"] = True
            self.cfg_ctrl["force_ctrl_method"] = "closed"
            self.cfg_ctrl["wrench_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.closed_loop_force.wrench_prop_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["force_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.closed_loop_force.force_ctrl_axes, device=self.device
            ).repeat((self.num_envs, 1))
        elif ctrl_type == "hybrid_force_motion":
            self.cfg_ctrl["motor_ctrl_mode"] = "manual"
            self.cfg_ctrl["gain_space"] = "task"
            self.cfg_ctrl["do_motion_ctrl"] = True
            self.cfg_ctrl["task_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.hybrid_force_motion.task_prop_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["task_deriv_gains"] = torch.tensor(
                self.cfg_task.ctrl.hybrid_force_motion.task_deriv_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_inertial_comp"] = True
            self.cfg_ctrl["motion_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.hybrid_force_motion.motion_ctrl_axes,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["do_force_ctrl"] = True
            self.cfg_ctrl["force_ctrl_method"] = "closed"
            self.cfg_ctrl["wrench_prop_gains"] = torch.tensor(
                self.cfg_task.ctrl.hybrid_force_motion.wrench_prop_gains,
                device=self.device,
            ).repeat((self.num_envs, 1))
            self.cfg_ctrl["force_ctrl_axes"] = torch.tensor(
                self.cfg_task.ctrl.hybrid_force_motion.force_ctrl_axes,
                device=self.device,
            ).repeat((self.num_envs, 1))

        if self.cfg_ctrl["motor_ctrl_mode"] == "gym":
            prop_gains = torch.cat(
                (
                    self.cfg_ctrl["joint_prop_gains"],
                    self.cfg_ctrl["gripper_prop_gains"],
                ),
                dim=-1,
            ).to("cpu")
            deriv_gains = torch.cat(
                (
                    self.cfg_ctrl["joint_deriv_gains"],
                    self.cfg_ctrl["gripper_deriv_gains"],
                ),
                dim=-1,
            ).to("cpu")
            # No tensor API for getting/setting actor DOF props; thus, loop required
            for env_ptr, franka_handle, prop_gain, deriv_gain in zip(
                self.env_ptrs, self.franka_handles, prop_gains, deriv_gains
            ):
                franka_dof_props = self.gym.get_actor_dof_properties(
                    env_ptr, franka_handle
                )
                franka_dof_props["driveMode"][:] = gymapi.DOF_MODE_POS
                # franka_dof_props["stiffness"] = prop_gain
                # franka_dof_props["damping"] = deriv_gain
                franka_dof_props["stiffness"][:] = 7000.0
                franka_dof_props["damping"][:] = 50.0
                self.gym.set_actor_dof_properties(
                    env_ptr, franka_handle, franka_dof_props
                )
        elif self.cfg_ctrl["motor_ctrl_mode"] == "manual":
            # No tensor API for getting/setting actor DOF props; thus, loop required
            for env_ptr, franka_handle in zip(self.env_ptrs, self.franka_handles):
                franka_dof_props = self.gym.get_actor_dof_properties(
                    env_ptr, franka_handle
                )
                franka_dof_props["driveMode"][:] = gymapi.DOF_MODE_EFFORT
                franka_dof_props["stiffness"][:] = 0.0  # zero passive stiffness
                franka_dof_props["damping"][:] = 0.0  # zero passive damping
                self.gym.set_actor_dof_properties(
                    env_ptr, franka_handle, franka_dof_props
                )

    def generate_replay_signals(self, joint_pos):
        ctrl_target_gripper_dof_pos = self.ctrl_target_dof_pos[:, -2:].clone()
        ctrl_target_gripper_dof_pos = torch.tensor([[0.025, 0.025]], device=self.device)
        self.ctrl_target_dof_pos = torch.zeros(
            (self.cfg_ctrl["num_envs"], 9), device=self.device
        )
        self.ctrl_target_dof_pos[:, 0:7] = joint_pos
        self.ctrl_target_dof_pos[:, 7:9] = ctrl_target_gripper_dof_pos
        self.gym.set_dof_position_target_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.ctrl_target_dof_pos),
            gymtorch.unwrap_tensor(self.franka_actor_ids_sim),
            len(self.franka_actor_ids_sim),
        )

    def generate_ctrl_signals(self):
        """Get Jacobian. Set Franka DOF position targets or DOF torques."""

        # Get desired Jacobian
        if self.cfg_ctrl["jacobian_type"] == "geometric":
            self.fingertip_midpoint_jacobian_tf = self.fingertip_midpoint_jacobian
        elif self.cfg_ctrl["jacobian_type"] == "analytic":
            self.fingertip_midpoint_jacobian_tf = fc.get_analytic_jacobian(
                fingertip_quat=self.fingertip_quat,
                fingertip_jacobian=self.fingertip_midpoint_jacobian,
                num_envs=self.num_envs,
                device=self.device,
            )

        # Set PD joint pos target or joint torque
        if self.cfg_ctrl["motor_ctrl_mode"] == "gym":
            self._set_dof_pos_target()
        elif self.cfg_ctrl["motor_ctrl_mode"] == "manual":
            self._set_dof_torque()

    def _set_dof_pos_target(self):
        """Set Franka DOF position target to move fingertips towards target pose."""

        self.ctrl_target_dof_pos = fc.compute_dof_pos_target(
            cfg_ctrl=self.cfg_ctrl,
            arm_dof_pos=self.arm_dof_pos,
            fingertip_midpoint_pos=self.fingertip_midpoint_pos,
            fingertip_midpoint_quat=self.fingertip_midpoint_quat,
            jacobian=self.fingertip_midpoint_jacobian_tf,
            ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos,
            ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat,
            ctrl_target_gripper_dof_pos=self.ctrl_target_gripper_dof_pos,
            device=self.device,
        )

        self.gym.set_dof_position_target_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.ctrl_target_dof_pos),
            gymtorch.unwrap_tensor(self.franka_actor_ids_sim),
            len(self.franka_actor_ids_sim),
        )

    def _set_dof_torque(self):
        """Set Franka DOF torque to move fingertips towards target pose."""

        self.dof_torque = fc.compute_dof_torque(
            cfg_ctrl=self.cfg_ctrl,
            dof_pos=self.dof_pos,
            dof_vel=self.dof_vel,
            fingertip_midpoint_pos=self.fingertip_midpoint_pos,
            fingertip_midpoint_quat=self.fingertip_midpoint_quat,
            fingertip_midpoint_linvel=self.fingertip_midpoint_linvel,
            fingertip_midpoint_angvel=self.fingertip_midpoint_angvel,
            left_finger_force=self.left_finger_force,
            right_finger_force=self.right_finger_force,
            jacobian=self.fingertip_midpoint_jacobian_tf,
            arm_mass_matrix=self.arm_mass_matrix,
            ctrl_target_gripper_dof_pos=self.ctrl_target_gripper_dof_pos,
            ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos,
            ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat,
            ctrl_target_fingertip_contact_wrench=self.ctrl_target_fingertip_contact_wrench,
            device=self.device,
        )

        self.gym.set_dof_actuation_force_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_torque),
            gymtorch.unwrap_tensor(self.franka_actor_ids_sim),
            len(self.franka_actor_ids_sim),
        )

    def print_sdf_warning(self):
        """Generate SDF warning message."""

        logger.warn(
            "Please be patient: SDFs may be generating, which may take a few minutes. Terminating prematurely may result in a corrupted SDF cache."
        )

    def enable_gravity(self, gravity_mag):
        """Enable gravity."""

        sim_params = self.gym.get_sim_params(self.sim)
        sim_params.gravity.z = -gravity_mag
        self.gym.set_sim_params(self.sim, sim_params)

    def disable_gravity(self):
        """Disable gravity."""

        sim_params = self.gym.get_sim_params(self.sim)
        sim_params.gravity.z = 0.0
        self.gym.set_sim_params(self.sim, sim_params)

    def export_scene(self, label):
        """Export scene to USD."""

        usd_export_options = gymapi.UsdExportOptions()
        usd_export_options.export_physics = False

        usd_exporter = self.gym.create_usd_exporter(usd_export_options)
        self.gym.export_usd_sim(usd_exporter, self.sim, label)
        sys.exit()

    def extract_poses(self):
        """Extract poses of all bodies."""

        if not hasattr(self, "export_pos"):
            self.export_pos = []
            self.export_rot = []
            self.frame_count = 0

        pos = self.body_pos
        rot = self.body_quat

        self.export_pos.append(pos.cpu().numpy().copy())
        self.export_rot.append(rot.cpu().numpy().copy())
        self.frame_count += 1

        if len(self.export_pos) == self.max_episode_length:
            output_dir = self.__class__.__name__
            save_dir = os.path.join("usd", output_dir)
            os.makedirs(output_dir, exist_ok=True)

            print(f"Exporting poses to {output_dir}...")
            np.save(
                os.path.join(save_dir, "body_position.npy"), np.array(self.export_pos)
            )
            np.save(
                os.path.join(save_dir, "body_rotation.npy"), np.array(self.export_rot)
            )
            print("Export completed.")
            sys.exit()

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_base.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: control module.

Imported by base, environment, and task classes. Not directly executed.
"""

import math
import torch

from isaacgymenvs.utils import torch_jit_utils as torch_utils



def compute_dof_pos_target(cfg_ctrl,
                           arm_dof_pos,
                           fingertip_midpoint_pos,
                           fingertip_midpoint_quat,
                           jacobian,
                           ctrl_target_fingertip_midpoint_pos,
                           ctrl_target_fingertip_midpoint_quat,
                           ctrl_target_gripper_dof_pos,
                           device):
    """Compute Franka DOF position target to move fingertips towards target pose."""

    ctrl_target_dof_pos = torch.zeros((cfg_ctrl['num_envs'], 9), device=device)

    pos_error, axis_angle_error = get_pose_error(
        fingertip_midpoint_pos=fingertip_midpoint_pos,
        fingertip_midpoint_quat=fingertip_midpoint_quat,
        ctrl_target_fingertip_midpoint_pos=ctrl_target_fingertip_midpoint_pos,
        ctrl_target_fingertip_midpoint_quat=ctrl_target_fingertip_midpoint_quat,
        jacobian_type=cfg_ctrl['jacobian_type'],
        rot_error_type='axis_angle')

    delta_fingertip_pose = torch.cat((pos_error, axis_angle_error), dim=1)
    delta_arm_dof_pos = _get_delta_dof_pos(delta_pose=delta_fingertip_pose,
                                           ik_method=cfg_ctrl['ik_method'],
                                           jacobian=jacobian,
                                           device=device)

    ctrl_target_dof_pos[:, 0:7] = arm_dof_pos + delta_arm_dof_pos
    ctrl_target_dof_pos[:, 7:9] = ctrl_target_gripper_dof_pos  # gripper finger joints

    return ctrl_target_dof_pos


def compute_dof_torque(cfg_ctrl,
                       dof_pos,
                       dof_vel,
                       fingertip_midpoint_pos,
                       fingertip_midpoint_quat,
                       fingertip_midpoint_linvel,
                       fingertip_midpoint_angvel,
                       left_finger_force,
                       right_finger_force,
                       jacobian,
                       arm_mass_matrix,
                       ctrl_target_gripper_dof_pos,
                       ctrl_target_fingertip_midpoint_pos,
                       ctrl_target_fingertip_midpoint_quat,
                       ctrl_target_fingertip_contact_wrench,
                       device):
    """Compute Franka DOF torque to move fingertips towards target pose."""
    # References:
    # 1) https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf
    # 2) Modern Robotics

    dof_torque = torch.zeros((cfg_ctrl['num_envs'], 9), device=device)

    if cfg_ctrl['gain_space'] == 'joint':
        pos_error, axis_angle_error = get_pose_error(
            fingertip_midpoint_pos=fingertip_midpoint_pos,
            fingertip_midpoint_quat=fingertip_midpoint_quat,
            ctrl_target_fingertip_midpoint_pos=ctrl_target_fingertip_midpoint_pos,
            ctrl_target_fingertip_midpoint_quat=ctrl_target_fingertip_midpoint_quat,
            jacobian_type=cfg_ctrl['jacobian_type'],
            rot_error_type='axis_angle')
        delta_fingertip_pose = torch.cat((pos_error, axis_angle_error), dim=1)

        # Set tau = k_p * joint_pos_error - k_d * joint_vel_error (ETH eq. 3.72)
        delta_arm_dof_pos = _get_delta_dof_pos(delta_pose=delta_fingertip_pose,
                                               ik_method=cfg_ctrl['ik_method'],
                                               jacobian=jacobian,
                                               device=device)
        dof_torque[:, 0:7] = cfg_ctrl['joint_prop_gains'] * delta_arm_dof_pos + \
                             cfg_ctrl['joint_deriv_gains'] * (0.0 - dof_vel[:, 0:7])

        if cfg_ctrl['do_inertial_comp']:
            # Set tau = M * tau, where M is the joint-space mass matrix
            arm_mass_matrix_joint = arm_mass_matrix
            dof_torque[:, 0:7] = (arm_mass_matrix_joint @ dof_torque[:, 0:7].unsqueeze(-1)).squeeze(-1)

    elif cfg_ctrl['gain_space'] == 'task':
        task_wrench = torch.zeros((cfg_ctrl['num_envs'], 6), device=device)

        if cfg_ctrl['do_motion_ctrl']:
            pos_error, axis_angle_error = get_pose_error(
                fingertip_midpoint_pos=fingertip_midpoint_pos,
                fingertip_midpoint_quat=fingertip_midpoint_quat,
                ctrl_target_fingertip_midpoint_pos=ctrl_target_fingertip_midpoint_pos,
                ctrl_target_fingertip_midpoint_quat=ctrl_target_fingertip_midpoint_quat,
                jacobian_type=cfg_ctrl['jacobian_type'],
                rot_error_type='axis_angle')
            delta_fingertip_pose = torch.cat((pos_error, axis_angle_error), dim=1)

            # Set tau = k_p * task_pos_error - k_d * task_vel_error (building towards eq. 3.96-3.98)
            task_wrench_motion = _apply_task_space_gains(delta_fingertip_pose=delta_fingertip_pose,
                                                         fingertip_midpoint_linvel=fingertip_midpoint_linvel,
                                                         fingertip_midpoint_angvel=fingertip_midpoint_angvel,
                                                         task_prop_gains=cfg_ctrl['task_prop_gains'],
                                                         task_deriv_gains=cfg_ctrl['task_deriv_gains'])

            if cfg_ctrl['do_inertial_comp']:
                # Set tau = Lambda * tau, where Lambda is the task-space mass matrix
                jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
                arm_mass_matrix_task = torch.inverse(jacobian @ torch.inverse(arm_mass_matrix) @ jacobian_T)  # ETH eq. 3.86; geometric Jacobian is assumed
                task_wrench_motion = (arm_mass_matrix_task @ task_wrench_motion.unsqueeze(-1)).squeeze(-1)

            task_wrench = task_wrench + torch.tensor(cfg_ctrl['motion_ctrl_axes'], device=device).unsqueeze(0) * task_wrench_motion

        if cfg_ctrl['do_force_ctrl']:
            # Set tau = tau + F_t, where F_t is the target contact wrench
            task_wrench_force = torch.zeros((cfg_ctrl['num_envs'], 6), device=device)
            task_wrench_force = task_wrench_force + ctrl_target_fingertip_contact_wrench  # open-loop force control (building towards ETH eq. 3.96-3.98)

            if cfg_ctrl['force_ctrl_method'] == 'closed':
                force_error, torque_error = _get_wrench_error(
                    left_finger_force=left_finger_force,
                    right_finger_force=right_finger_force,
                    ctrl_target_fingertip_contact_wrench=ctrl_target_fingertip_contact_wrench,
                    num_envs=cfg_ctrl['num_envs'],
                    device=device)

                # Set tau = tau + k_p * contact_wrench_error
                task_wrench_force = task_wrench_force + cfg_ctrl['wrench_prop_gains'] * torch.cat(
                    (force_error, torque_error), dim=1)  # part of Modern Robotics eq. 11.61

            task_wrench = task_wrench + torch.tensor(cfg_ctrl['force_ctrl_axes'], device=device).unsqueeze(
                0) * task_wrench_force

        # Set tau = J^T * tau, i.e., map tau into joint space as desired
        jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
        dof_torque[:, 0:7] = (jacobian_T @ task_wrench.unsqueeze(-1)).squeeze(-1)

    dof_torque[:, 7:9] = cfg_ctrl['gripper_prop_gains'] * (ctrl_target_gripper_dof_pos - dof_pos[:, 7:9]) + \
                         cfg_ctrl['gripper_deriv_gains'] * (0.0 - dof_vel[:, 7:9])  # gripper finger joints
    dof_torque = torch.clamp(dof_torque, min=-100.0, max=100.0)

    return dof_torque


def get_pose_error(fingertip_midpoint_pos,
                   fingertip_midpoint_quat,
                   ctrl_target_fingertip_midpoint_pos,
                   ctrl_target_fingertip_midpoint_quat,
                   jacobian_type,
                   rot_error_type):
    """Compute task-space error between target Franka fingertip pose and current pose."""
    # Reference: https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf

    # Compute pos error
    pos_error = ctrl_target_fingertip_midpoint_pos - fingertip_midpoint_pos

    # Compute rot error
    if jacobian_type == 'geometric':  # See example 2.9.8; note use of J_g and transformation between rotation vectors
        # Compute quat error (i.e., difference quat)
        # Reference: https://personal.utdallas.edu/~sxb027100/dock/quat.html
        fingertip_midpoint_quat_norm = torch_utils.quat_mul(fingertip_midpoint_quat,
                                                            torch_utils.quat_conjugate(fingertip_midpoint_quat))[:, 3]  # scalar component
        fingertip_midpoint_quat_inv = torch_utils.quat_conjugate(
            fingertip_midpoint_quat) / fingertip_midpoint_quat_norm.unsqueeze(-1)
        quat_error = torch_utils.quat_mul(ctrl_target_fingertip_midpoint_quat, fingertip_midpoint_quat_inv)

        # Convert to axis-angle error
        axis_angle_error = axis_angle_from_quat(quat_error)

    elif jacobian_type == 'analytic':  # See example 2.9.7; note use of J_a and difference of rotation vectors
        # Compute axis-angle error
        axis_angle_error = axis_angle_from_quat(ctrl_target_fingertip_midpoint_quat)\
                           - axis_angle_from_quat(fingertip_midpoint_quat)

    if rot_error_type == 'quat':
        return pos_error, quat_error
    elif rot_error_type == 'axis_angle':
        return pos_error, axis_angle_error


def _get_wrench_error(left_finger_force,
                      right_finger_force,
                      ctrl_target_fingertip_contact_wrench,
                      num_envs,
                      device):
    """Compute task-space error between target Franka fingertip contact wrench and current wrench."""

    fingertip_contact_wrench = torch.zeros((num_envs, 6), device=device)

    fingertip_contact_wrench[:, 0:3] = left_finger_force + right_finger_force  # net contact force on fingers
    # Cols 3 to 6 are all zeros, as we do not have enough information

    force_error = ctrl_target_fingertip_contact_wrench[:, 0:3] - (-fingertip_contact_wrench[:, 0:3])
    torque_error = ctrl_target_fingertip_contact_wrench[:, 3:6] - (-fingertip_contact_wrench[:, 3:6])

    return force_error, torque_error


def _get_delta_dof_pos(delta_pose, ik_method, jacobian, device):
    """Get delta Franka DOF position from delta pose using specified IK method."""
    # References:
    # 1) https://www.cs.cmu.edu/~15464-s13/lectures/lecture6/iksurvey.pdf
    # 2) https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf (p. 47)

    if ik_method == 'pinv':  # Jacobian pseudoinverse
        k_val = 1.0
        jacobian_pinv = torch.linalg.pinv(jacobian)
        delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
        delta_dof_pos = delta_dof_pos.squeeze(-1)

    elif ik_method == 'trans':  # Jacobian transpose
        k_val = 1.0
        jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
        delta_dof_pos = k_val * jacobian_T @ delta_pose.unsqueeze(-1)
        delta_dof_pos = delta_dof_pos.squeeze(-1)

    elif ik_method == 'dls':  # damped least squares (Levenberg-Marquardt)
        lambda_val = 0.1
        jacobian_T = torch.transpose(jacobian, dim0=1, dim1=2)
        lambda_matrix = (lambda_val ** 2) * torch.eye(n=jacobian.shape[1], device=device)
        delta_dof_pos = jacobian_T @ torch.inverse(jacobian @ jacobian_T + lambda_matrix) @ delta_pose.unsqueeze(-1)
        delta_dof_pos = delta_dof_pos.squeeze(-1)

    elif ik_method == 'svd':  # adaptive SVD
        k_val = 1.0
        U, S, Vh = torch.linalg.svd(jacobian)
        S_inv = 1. / S
        min_singular_value = 1.0e-5
        S_inv = torch.where(S > min_singular_value, S_inv, torch.zeros_like(S_inv))
        jacobian_pinv = torch.transpose(Vh, dim0=1, dim1=2)[:, :, :6] @ torch.diag_embed(S_inv) @ torch.transpose(U, dim0=1, dim1=2)
        delta_dof_pos = k_val * jacobian_pinv @ delta_pose.unsqueeze(-1)
        delta_dof_pos = delta_dof_pos.squeeze(-1)

    return delta_dof_pos


def _apply_task_space_gains(delta_fingertip_pose,
                            fingertip_midpoint_linvel,
                            fingertip_midpoint_angvel,
                            task_prop_gains,
                            task_deriv_gains):
    """Interpret PD gains as task-space gains. Apply to task-space error."""

    task_wrench = torch.zeros_like(delta_fingertip_pose)

    # Apply gains to lin error components
    lin_error = delta_fingertip_pose[:, 0:3]
    task_wrench[:, 0:3] = task_prop_gains[:, 0:3] * lin_error + \
                          task_deriv_gains[:, 0:3] * (0.0 - fingertip_midpoint_linvel)

    # Apply gains to rot error components
    rot_error = delta_fingertip_pose[:, 3:6]
    task_wrench[:, 3:6] = task_prop_gains[:, 3:6] * rot_error + \
                          task_deriv_gains[:, 3:6] * (0.0 - fingertip_midpoint_angvel)

    return task_wrench


def get_analytic_jacobian(fingertip_quat, fingertip_jacobian, num_envs, device):
    """Convert geometric Jacobian to analytic Jacobian."""
    # Reference: https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2018/RD_HS2018script.pdf
    # NOTE: Gym returns world-space geometric Jacobians by default

    batch = num_envs

    # Overview:
    # x = [x_p; x_r]
    # From eq. 2.189 and 2.192, x_dot = J_a @ q_dot = (E_inv @ J_g) @ q_dot
    # From eq. 2.191, E = block(E_p, E_r); thus, E_inv = block(E_p_inv, E_r_inv)
    # Eq. 2.12 gives an expression for E_p_inv
    # Eq. 2.107 gives an expression for E_r_inv

    # Compute E_inv_top (i.e., [E_p_inv, 0])
    I = torch.eye(3, device=device)
    E_p_inv = I.repeat((batch, 1)).reshape(batch, 3, 3)
    E_inv_top = torch.cat((E_p_inv, torch.zeros((batch, 3, 3), device=device)), dim=2)

    # Compute E_inv_bottom (i.e., [0, E_r_inv])
    fingertip_axis_angle = axis_angle_from_quat(fingertip_quat)
    fingertip_axis_angle_cross = get_skew_symm_matrix(fingertip_axis_angle, device=device)
    fingertip_angle = torch.linalg.vector_norm(fingertip_axis_angle, dim=1)
    factor_1 = 1 / (fingertip_angle ** 2)
    factor_2 = 1 - fingertip_angle * 0.5 * torch.sin(fingertip_angle) / (1 - torch.cos(fingertip_angle))
    factor_3 = factor_1 * factor_2
    E_r_inv = I \
              - 1 * 0.5 * fingertip_axis_angle_cross \
              + (fingertip_axis_angle_cross @ fingertip_axis_angle_cross) * factor_3.unsqueeze(-1).repeat((1, 3 * 3)).reshape((batch, 3, 3))
    E_inv_bottom = torch.cat((torch.zeros((batch, 3, 3), device=device), E_r_inv), dim=2)

    E_inv = torch.cat((E_inv_top.reshape((batch, 3 * 6)), E_inv_bottom.reshape((batch, 3 * 6))), dim=1).reshape((batch, 6, 6))

    J_a = E_inv @ fingertip_jacobian

    return J_a


def get_skew_symm_matrix(vec, device):
    """Convert vector to skew-symmetric matrix."""
    # Reference: https://en.wikipedia.org/wiki/Cross_product#Conversion_to_matrix_multiplication

    batch = vec.shape[0]
    I = torch.eye(3, device=device)
    skew_symm = torch.transpose(torch.cross(vec.repeat((1, 3)).reshape((batch * 3, 3)),
                                            I.repeat((batch, 1)))
                                .reshape(batch, 3, 3),
                                dim0=1,
                                dim1=2)

    return skew_symm


def translate_along_local_z(pos, quat, offset, device):
    """Translate global body position along local Z-axis and express in global coordinates."""

    num_vecs = pos.shape[0]
    offset_vec = offset * torch.tensor([0.0, 0.0, 1.0], device=device).repeat((num_vecs, 1))
    _, translated_pos = torch_utils.tf_combine(q1=quat,
                                               t1=pos,
                                               q2=torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).repeat((num_vecs, 1)),
                                               t2=offset_vec)

    return translated_pos


def axis_angle_from_euler(euler):
    """Convert tensor of Euler angles to tensor of axis-angles."""

    quat = torch_utils.quat_from_euler_xyz(roll=euler[:, 0], pitch=euler[:, 1], yaw=euler[:, 2])
    quat = quat * torch.sign(quat[:, 3]).unsqueeze(-1)  # smaller rotation 
    axis_angle = axis_angle_from_quat(quat)

    return axis_angle


def axis_angle_from_quat(quat, eps=1.0e-6):
    """Convert tensor of quaternions to tensor of axis-angles."""
    # Reference: https://github.com/facebookresearch/pytorch3d/blob/bee31c48d3d36a8ea268f9835663c52ff4a476ec/pytorch3d/transforms/rotation_conversions.py#L516-L544

    mag = torch.linalg.norm(quat[:, 0:3], dim=1)
    half_angle = torch.atan2(mag, quat[:, 3])
    angle = 2.0 * half_angle
    sin_half_angle_over_angle = torch.where(torch.abs(angle) > eps,
                                            torch.sin(half_angle) / angle,
                                            1 / 2 - angle ** 2.0 / 48)
    axis_angle = quat[:, 0:3] / sin_half_angle_over_angle.unsqueeze(-1)

    return axis_angle


def axis_angle_from_quat_naive(quat):
    """Convert tensor of quaternions to tensor of axis-angles."""
    # Reference: https://en.wikipedia.org/wiki/quats_and_spatial_rotation#Recovering_the_axis-angle_representation
    # NOTE: Susceptible to undesirable behavior due to divide-by-zero

    mag = torch.linalg.vector_norm(quat[:, 0:3], dim=1)  # zero when quat = [0, 0, 0, 1]
    axis = quat[:, 0:3] / mag.unsqueeze(-1)
    angle = 2.0 * torch.atan2(mag, quat[:, 3])
    axis_angle = axis * angle.unsqueeze(-1)

    return axis_angle


def get_rand_quat(num_quats, device):
    """Generate tensor of random quaternions."""
    # Reference: http://planning.cs.uiuc.edu/node198.html

    u = torch.rand((num_quats, 3), device=device)
    quat = torch.zeros((num_quats, 4), device=device)
    quat[:, 0] = torch.sqrt(1 - u[:, 0]) * torch.sin(2 * math.pi * u[:, 1])
    quat[:, 1] = torch.sqrt(1 - u[:, 0]) * torch.cos(2 * math.pi * u[:, 1])
    quat[:, 2] = torch.sqrt(u[:, 0]) * torch.sin(2 * math.pi * u[:, 2])
    quat[:, 3] = torch.sqrt(u[:, 0]) * torch.cos(2 * math.pi * u[:, 2])

    return quat


def get_nonrand_quat(num_quats, rot_perturbation, device):
    """Generate tensor of non-random quaternions by composing random Euler rotations."""

    quat = torch_utils.quat_from_euler_xyz(
        torch.rand((num_quats, 1), device=device).squeeze() * rot_perturbation * 2.0 - rot_perturbation,
        torch.rand((num_quats, 1), device=device).squeeze() * rot_perturbation * 2.0 - rot_perturbation,
        torch.rand((num_quats, 1), device=device).squeeze() * rot_perturbation * 2.0 - rot_perturbation)

    return quat

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_control.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: class for nut-bolt env.

Inherits base class and abstract environment class. Inherited by nut-bolt task classes. Not directly executed.

Configuration defined in FactoryEnvNutBolt.yaml. Asset info defined in factory_asset_info_nut_bolt.yaml.
"""

import cv2
import hydra
import numpy as np
import os
import torch

from isaacgym import gymapi
from isaacgymenvs.tasks.ncf_tacto.factory_base import FactoryBase
import isaacgymenvs.tasks.ncf_tacto.factory_control as fc
from isaacgymenvs.tasks.ncf_tacto.factory_schema_class_env import FactoryABCEnv
from isaacgymenvs.tasks.ncf_tacto.factory_schema_config_env import (
    FactorySchemaConfigEnv,
)

from isaacgymenvs.tasks.ncf_tacto.digit_renderer import digit_renderer
from isaacgymenvs.tasks.ncf_tacto.pose import (
    xyzquat_to_tf_numpy,
    euler_angles_to_matrix,
)
from dataclasses import dataclass
import matplotlib.pyplot as plt

from multiprocessing import Process, Queue, Manager


@dataclass
class Cfg_Tacto:
    pixmm: float = 0.03
    width: int = 240
    height: int = 320
    cam_dist: float = 0.022
    shear_mag: float = 0.0
    pen_min: float = 0.0005
    pen_max: float = 0.002
    randomize: bool = False
    bg_id: int = 10


class FactoryEnvNutBolt(FactoryBase, FactoryABCEnv):
    def __init__(
        self,
        cfg,
        rl_device,
        sim_device,
        graphics_device_id,
        headless,
        virtual_screen_capture,
        force_render,
    ):
        """Initialize instance variables. Initialize environment superclass. Acquire tensors."""
        self.t_digits = 0
        self.subsample_digits = 5  # 5
        self._get_env_yaml_params()

        super().__init__(
            cfg,
            rl_device,
            sim_device,
            graphics_device_id,
            headless,
            virtual_screen_capture,
            force_render,
        )

        self.acquire_base_tensors()  # defined in superclass
        self._acquire_env_tensors()
        self.refresh_base_tensors()  # defined in superclass
        self.refresh_env_tensors()
        if cfg["sim"]["sim_digits"]:
            self.refresh_digit_sensors()

    def _get_env_yaml_params(self):
        """Initialize instance variables from YAML files."""

        cs = hydra.core.config_store.ConfigStore.instance()
        cs.store(name="factory_schema_config_env", node=FactorySchemaConfigEnv)

        config_path = (
            "task/NCFEnvNutBolt.yaml"  # relative to Hydra search path (cfg dir)
        )
        self.cfg_env = hydra.compose(config_name=config_path)
        self.cfg_env = self.cfg_env["task"]  # strip superfluous nesting

        asset_info_path = "../../assets/factory/yaml/factory_asset_info_nut_bolt.yaml"
        self.asset_info_nut_bolt = hydra.compose(config_name=asset_info_path)
        self.asset_info_nut_bolt = self.asset_info_nut_bolt[""][""][""][""][""][""][
            "assets"
        ]["factory"][
            "yaml"
        ]  # strip superfluous nesting

    def create_envs(self):
        """Set env options. Import assets. Create actors."""

        lower = gymapi.Vec3(
            -self.cfg_base.env.env_spacing, -self.cfg_base.env.env_spacing, 0.0
        )
        upper = gymapi.Vec3(
            self.cfg_base.env.env_spacing,
            self.cfg_base.env.env_spacing,
            self.cfg_base.env.env_spacing,
        )
        num_per_row = int(np.sqrt(self.num_envs))

        self.print_sdf_warning()
        franka_asset, table_asset = self.import_franka_assets()
        (
            nut_asset,
            bolt_asset,
            nut_mesh_file,
            nut_mesh_scale,
            nut_pointcloud_file,
        ) = self._import_env_assets()

        # self.init_digit_render(object_mesh=nut_mesh_file, object_mesh_scale=nut_mesh_scale)

        self._create_actors(
            lower,
            upper,
            num_per_row,
            franka_asset,
            nut_asset,
            bolt_asset,
            table_asset,
            nut_mesh_file,
            nut_mesh_scale,
            nut_pointcloud_file,
        )

    def _import_env_assets(self):
        """Set nut and bolt asset options. Import assets."""

        urdf_root = os.path.join(
            os.path.dirname(__file__), "..", "..", "..", "assets", "factory", "urdf"
        )
        mesh_root = os.path.join(
            os.path.dirname(__file__), "..", "..", "..", "assets", "factory", "mesh"
        )

        nut_options = gymapi.AssetOptions()
        nut_options.flip_visual_attachments = False
        nut_options.fix_base_link = False
        nut_options.thickness = 0.0  # default = 0.02
        nut_options.armature = 0.0  # default = 0.0
        nut_options.use_physx_armature = True
        nut_options.linear_damping = 0.0  # default = 0.0
        nut_options.max_linear_velocity = 1000.0  # default = 1000.0
        nut_options.angular_damping = 0.0  # default = 0.5
        nut_options.max_angular_velocity = 64.0  # default = 64.0
        nut_options.disable_gravity = False
        nut_options.enable_gyroscopic_forces = True
        nut_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        nut_options.use_mesh_materials = False
        if self.cfg_base.mode.export_scene:
            nut_options.mesh_normal_mode = gymapi.COMPUTE_PER_FACE

        bolt_options = gymapi.AssetOptions()
        bolt_options.flip_visual_attachments = False
        bolt_options.fix_base_link = True
        bolt_options.thickness = 0.0  # default = 0.02
        bolt_options.armature = 0.0  # default = 0.0
        bolt_options.use_physx_armature = True
        bolt_options.linear_damping = 0.0  # default = 0.0
        bolt_options.max_linear_velocity = 1000.0  # default = 1000.0
        bolt_options.angular_damping = 0.0  # default = 0.5
        bolt_options.max_angular_velocity = 64.0  # default = 64.0
        bolt_options.disable_gravity = False
        bolt_options.enable_gyroscopic_forces = True
        bolt_options.default_dof_drive_mode = gymapi.DOF_MODE_NONE
        bolt_options.use_mesh_materials = False
        if self.cfg_base.mode.export_scene:
            bolt_options.mesh_normal_mode = gymapi.COMPUTE_PER_FACE

        nut_assets = []
        bolt_assets = []
        for subassembly in self.cfg_env.env.desired_subassemblies:
            components = list(self.asset_info_nut_bolt[subassembly])
            nut_file = (
                self.asset_info_nut_bolt[subassembly][components[0]]["urdf_path"]
                + ".urdf"
            )
            bolt_file = (
                self.asset_info_nut_bolt[subassembly][components[1]]["urdf_path"]
                + ".urdf"
            )
            nut_options.density = self.cfg_env.env.nut_bolt_density
            bolt_options.density = self.cfg_env.env.nut_bolt_density
            nut_asset = self.gym.load_asset(self.sim, urdf_root, nut_file, nut_options)
            bolt_asset = self.gym.load_asset(
                self.sim, urdf_root, bolt_file, bolt_options
            )
            nut_assets.append(nut_asset)
            bolt_assets.append(bolt_asset)

        nut_mesh_file = (
            mesh_root
            + "/"
            + self.asset_info_nut_bolt[subassembly][components[0]]["mesh_path"]
        )
        nut_mesh_scale = self.asset_info_nut_bolt[subassembly][components[0]][
            "mesh_scale"
        ]

        nut_pointcloud_file = (
            mesh_root
            + "/"
            + self.asset_info_nut_bolt[subassembly][components[0]]["pointcloud_path"]
        )

        return (
            nut_assets,
            bolt_assets,
            nut_mesh_file,
            nut_mesh_scale,
            nut_pointcloud_file,
        )

    def _create_actors(
        self,
        lower,
        upper,
        num_per_row,
        franka_asset,
        nut_assets,
        bolt_assets,
        table_asset,
        nut_mesh_file,
        nut_mesh_scale,
        nut_pointcloud_file,
    ):
        """Set initial actor poses. Create actors. Set shape and DOF properties."""

        franka_pose = gymapi.Transform()
        franka_pose.p.x = self.cfg_base.env.franka_depth
        franka_pose.p.y = 0.0
        franka_pose.p.z = 0.0
        franka_pose.r = gymapi.Quat(0.0, 0.0, 1.0, 0.0)

        table_pose = gymapi.Transform()
        table_pose.p.x = 0.0
        table_pose.p.y = 0.0
        table_pose.p.z = self.cfg_base.env.table_height * 0.5
        table_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)

        self.env_ptrs = []
        self.franka_handles = []
        self.nut_handles = []
        self.bolt_handles = []
        self.table_handles = []
        self.shape_ids = []
        self.franka_actor_ids_sim = []  # within-sim indices
        self.nut_actor_ids_sim = []  # within-sim indices
        self.bolt_actor_ids_sim = []  # within-sim indices
        self.table_actor_ids_sim = []  # within-sim indices
        actor_count = 0

        self.nut_heights = []
        self.nut_widths_max = []
        self.bolt_widths = []
        self.bolt_head_heights = []
        self.bolt_shank_lengths = []
        self.thread_pitches = []

        self.nut_mesh_scale = nut_mesh_scale
        self.nut_mesh_file = nut_mesh_file
        nut_pc = np.load(nut_pointcloud_file)
        self.nut_pointcloud = torch.tensor(nut_pc, dtype=torch.float32).to(self.device)

        print(f"[isaacgym] creating {self.num_envs} envs with digit sensors")

        for i in range(self.num_envs):
            env_ptr = self.gym.create_env(self.sim, lower, upper, num_per_row)

            if self.cfg_env.sim.disable_franka_collisions:
                franka_handle = self.gym.create_actor(
                    env_ptr,
                    franka_asset,
                    franka_pose,
                    "franka",
                    i + self.num_envs,
                    0,
                    0,
                )
            else:
                franka_handle = self.gym.create_actor(
                    env_ptr, franka_asset, franka_pose, "franka", i, 0, 0
                )
            self.franka_actor_ids_sim.append(actor_count)
            actor_count += 1

            j = np.random.randint(0, len(self.cfg_env.env.desired_subassemblies))
            subassembly = self.cfg_env.env.desired_subassemblies[j]
            components = list(self.asset_info_nut_bolt[subassembly])

            nut_pose = gymapi.Transform()
            nut_pose.p.x = 0.0
            nut_pose.p.y = self.cfg_env.env.nut_lateral_offset
            nut_pose.p.z = self.cfg_base.env.table_height
            nut_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)
            nut_color = gymapi.Vec3(0.4, 0.0, 0.1)
            nut_handle = self.gym.create_actor(
                env_ptr, nut_assets[j], nut_pose, "nut", i, 0, 0
            )
            self.gym.set_rigid_body_color(
                env_ptr, nut_handle, 0, gymapi.MESH_VISUAL, nut_color
            )
            self.nut_actor_ids_sim.append(actor_count)
            actor_count += 1

            nut_height = self.asset_info_nut_bolt[subassembly][components[0]]["height"]
            nut_width_max = self.asset_info_nut_bolt[subassembly][components[0]][
                "width_max"
            ]
            self.nut_heights.append(nut_height)
            self.nut_widths_max.append(nut_width_max)

            bolt_pose = gymapi.Transform()
            bolt_pose.p.x = 0.0
            bolt_pose.p.y = 0.0
            bolt_pose.p.z = self.cfg_base.env.table_height
            bolt_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)
            bolt_color = gymapi.Vec3(0.1, 0.6, 0.1)
            bolt_handle = self.gym.create_actor(
                env_ptr, bolt_assets[j], bolt_pose, "bolt", i, 0, 0
            )
            self.gym.set_rigid_body_color(
                env_ptr, bolt_handle, 0, gymapi.MESH_VISUAL, bolt_color
            )
            self.bolt_actor_ids_sim.append(actor_count)
            actor_count += 1

            bolt_width = self.asset_info_nut_bolt[subassembly][components[1]]["width"]
            bolt_head_height = self.asset_info_nut_bolt[subassembly][components[1]][
                "head_height"
            ]
            bolt_shank_length = self.asset_info_nut_bolt[subassembly][components[1]][
                "shank_length"
            ]
            self.bolt_widths.append(bolt_width)
            self.bolt_head_heights.append(bolt_head_height)
            self.bolt_shank_lengths.append(bolt_shank_length)

            thread_pitch = self.asset_info_nut_bolt[subassembly]["thread_pitch"]
            self.thread_pitches.append(thread_pitch)

            table_handle = self.gym.create_actor(
                env_ptr, table_asset, table_pose, "table", i, 0, 0
            )
            self.table_actor_ids_sim.append(actor_count)
            actor_count += 1

            link7_id = self.gym.find_actor_rigid_body_index(
                env_ptr, franka_handle, "panda_link7", gymapi.DOMAIN_ACTOR
            )
            hand_id = self.gym.find_actor_rigid_body_index(
                env_ptr, franka_handle, "panda_hand", gymapi.DOMAIN_ACTOR
            )
            left_finger_id = self.gym.find_actor_rigid_body_index(
                env_ptr, franka_handle, "panda_leftfinger", gymapi.DOMAIN_ACTOR
            )
            right_finger_id = self.gym.find_actor_rigid_body_index(
                env_ptr, franka_handle, "panda_rightfinger", gymapi.DOMAIN_ACTOR
            )
            self.shape_ids = [link7_id, hand_id, left_finger_id, right_finger_id]

            franka_shape_props = self.gym.get_actor_rigid_shape_properties(
                env_ptr, franka_handle
            )
            franka_shape_index = self.gym.get_actor_rigid_body_shape_indices(
                env_ptr, franka_handle
            )
            for shape_idx in self.shape_ids:
                shape_id = franka_shape_index[shape_idx].start
                franka_shape_props[
                    shape_id
                ].friction = self.cfg_base.env.franka_friction
                franka_shape_props[shape_id].rolling_friction = 0.0  # default = 0.0
                franka_shape_props[shape_id].torsion_friction = 0.0  # default = 0.0
                franka_shape_props[shape_id].restitution = 0.0  # default = 0.0
                franka_shape_props[shape_id].compliance = 0.0  # default = 0.0
                franka_shape_props[shape_id].thickness = 0.0  # default = 0.0
            self.gym.set_actor_rigid_shape_properties(
                env_ptr, franka_handle, franka_shape_props
            )

            nut_shape_props = self.gym.get_actor_rigid_shape_properties(
                env_ptr, nut_handle
            )
            nut_shape_props[0].friction = self.cfg_env.env.nut_bolt_friction
            nut_shape_props[0].rolling_friction = 0.0  # default = 0.0
            nut_shape_props[0].torsion_friction = 0.0  # default = 0.0
            nut_shape_props[0].restitution = 0.0  # default = 0.0
            nut_shape_props[0].compliance = 0.0  # default = 0.0
            nut_shape_props[0].thickness = 0.0  # default = 0.0
            self.gym.set_actor_rigid_shape_properties(
                env_ptr, nut_handle, nut_shape_props
            )

            bolt_shape_props = self.gym.get_actor_rigid_shape_properties(
                env_ptr, bolt_handle
            )
            bolt_shape_props[0].friction = self.cfg_env.env.nut_bolt_friction
            bolt_shape_props[0].rolling_friction = 0.0  # default = 0.0
            bolt_shape_props[0].torsion_friction = 0.0  # default = 0.0
            bolt_shape_props[0].restitution = 0.0  # default = 0.0
            bolt_shape_props[0].compliance = 0.0  # default = 0.0
            bolt_shape_props[0].thickness = 0.0  # default = 0.0
            self.gym.set_actor_rigid_shape_properties(
                env_ptr, bolt_handle, bolt_shape_props
            )

            table_shape_props = self.gym.get_actor_rigid_shape_properties(
                env_ptr, table_handle
            )
            table_color = gymapi.Vec3(0.0, 0.0, 0.0)
            self.gym.set_rigid_body_color(
                env_ptr, table_handle, 0, gymapi.MESH_VISUAL, table_color
            )

            table_shape_props[0].friction = self.cfg_base.env.table_friction
            table_shape_props[0].rolling_friction = 0.0  # default = 0.0
            table_shape_props[0].torsion_friction = 0.0  # default = 0.0
            table_shape_props[0].restitution = 0.0  # default = 0.0
            table_shape_props[0].compliance = 0.0  # default = 0.0
            table_shape_props[0].thickness = 0.0  # default = 0.0
            self.gym.set_actor_rigid_shape_properties(
                env_ptr, table_handle, table_shape_props
            )

            self.franka_num_dofs = self.gym.get_actor_dof_count(env_ptr, franka_handle)

            self.gym.enable_actor_dof_force_sensors(env_ptr, franka_handle)

            self.env_ptrs.append(env_ptr)
            self.franka_handles.append(franka_handle)
            self.nut_handles.append(nut_handle)
            self.bolt_handles.append(bolt_handle)
            self.table_handles.append(table_handle)

        print(f"[isaacgym] environment creation complete")

        self.num_actors = int(actor_count / self.num_envs)  # per env
        self.num_bodies = self.gym.get_env_rigid_body_count(env_ptr)  # per env
        self.num_dofs = self.gym.get_env_dof_count(env_ptr)  # per env

        # For setting targets
        self.franka_actor_ids_sim = torch.tensor(
            self.franka_actor_ids_sim, dtype=torch.int32, device=self.device
        )
        self.nut_actor_ids_sim = torch.tensor(
            self.nut_actor_ids_sim, dtype=torch.int32, device=self.device
        )
        self.bolt_actor_ids_sim = torch.tensor(
            self.bolt_actor_ids_sim, dtype=torch.int32, device=self.device
        )

        # For extracting root pos/quat
        self.nut_actor_id_env = self.gym.find_actor_index(
            env_ptr, "nut", gymapi.DOMAIN_ENV
        )
        self.bolt_actor_id_env = self.gym.find_actor_index(
            env_ptr, "bolt", gymapi.DOMAIN_ENV
        )

        # For extracting body pos/quat, force, and Jacobian
        self.nut_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, nut_handle, "nut", gymapi.DOMAIN_ENV
        )
        self.bolt_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, bolt_handle, "bolt", gymapi.DOMAIN_ENV
        )
        self.hand_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_hand", gymapi.DOMAIN_ENV
        )
        self.left_finger_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_leftfinger", gymapi.DOMAIN_ENV
        )
        self.right_finger_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_rightfinger", gymapi.DOMAIN_ENV
        )
        self.fingertip_centered_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_fingertip_centered", gymapi.DOMAIN_ENV
        )
        self.left_fingertip_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_leftfinger_tip", gymapi.DOMAIN_ENV
        )
        self.right_fingertip_body_id_env = self.gym.find_actor_rigid_body_index(
            env_ptr, franka_handle, "panda_rightfinger_tip", gymapi.DOMAIN_ENV
        )

        # For computing body COM pos
        self.nut_heights = torch.tensor(self.nut_heights, device=self.device).unsqueeze(
            -1
        )
        self.bolt_head_heights = torch.tensor(
            self.bolt_head_heights, device=self.device
        ).unsqueeze(-1)

        # For setting initial state
        self.nut_widths_max = torch.tensor(
            self.nut_widths_max, device=self.device
        ).unsqueeze(-1)
        self.bolt_shank_lengths = torch.tensor(
            self.bolt_shank_lengths, device=self.device
        ).unsqueeze(-1)

        # For defining success or failure
        self.bolt_widths = torch.tensor(self.bolt_widths, device=self.device).unsqueeze(
            -1
        )
        self.thread_pitches = torch.tensor(
            self.thread_pitches, device=self.device
        ).unsqueeze(-1)

    def _acquire_env_tensors(self):
        """Acquire and wrap tensors. Create views."""

        self.nut_pos = self.root_pos[:, self.nut_actor_id_env, 0:3]
        self.nut_quat = self.root_quat[:, self.nut_actor_id_env, 0:4]
        self.nut_linvel = self.root_linvel[:, self.nut_actor_id_env, 0:3]
        self.nut_angvel = self.root_angvel[:, self.nut_actor_id_env, 0:3]

        self.bolt_pos = self.root_pos[:, self.bolt_actor_id_env, 0:3]
        self.bolt_quat = self.root_quat[:, self.bolt_actor_id_env, 0:4]

        self.nut_force = self.contact_force[:, self.nut_body_id_env, 0:3]

        self.bolt_force = self.contact_force[:, self.bolt_body_id_env, 0:3]

        self.nut_com_pos = fc.translate_along_local_z(
            pos=self.nut_pos,
            quat=self.nut_quat,
            offset=self.bolt_head_heights + self.nut_heights * 0.5,
            device=self.device,
        )
        self.nut_com_quat = self.nut_quat  # always equal
        self.nut_com_linvel = self.nut_linvel + torch.cross(
            self.nut_angvel, (self.nut_com_pos - self.nut_pos), dim=1
        )
        self.nut_com_angvel = self.nut_angvel  # always equal

    def refresh_env_tensors(self):
        """Refresh tensors."""
        # NOTE: Tensor refresh functions should be called once per step, before setters.

        self.nut_com_pos = fc.translate_along_local_z(
            pos=self.nut_pos,
            quat=self.nut_quat,
            offset=self.bolt_head_heights + self.nut_heights * 0.5,
            device=self.device,
        )
        self.nut_com_linvel = self.nut_linvel + torch.cross(
            self.nut_angvel, (self.nut_com_pos - self.nut_pos), dim=1
        )

    def render_digit(
        self, offset, nut_poses, left_finger_poses, right_finger_poses, queue
    ):
        sz = self.cfg["ncf"]["digits_size"]

        cfg_tacto = Cfg_Tacto()
        # add digit handle
        self.digit_render_left = digit_renderer(
            cfg=cfg_tacto,
            obj_path=self.nut_mesh_file,
            obj_quat=np.array([0.0, 0.0, 0.0, 1.0]),
            obj_scale=self.nut_mesh_scale,
            bg_id=np.random.randint(0, 25),
            headless=True,
        )

        self.digit_render_right = digit_renderer(
            cfg=cfg_tacto,
            obj_path=self.nut_mesh_file,
            obj_quat=np.array([0.0, 0.0, 0.0, 1.0]),
            obj_scale=self.nut_mesh_scale,
            bg_id=np.random.randint(0, 25),
            headless=True,
        )

        bg_left = self.digit_render_left.renderer.background
        bg_right = self.digit_render_right.renderer.background

        lf = np.zeros((len(nut_poses), sz, sz, 3))
        rf = np.zeros((len(nut_poses), sz, sz, 3))

        for i in range(len(nut_poses)):
            self.digit_render_left.update_object_pose(nut_poses[i])
            self.digit_render_left.update_pose_given_pose(left_finger_poses[i])
            im_lf, hm_lf, mask_lf = self.digit_render_left.render()
            im_lf = self._subtract_bg(im_lf, bg_left)
            im_lf = cv2.resize(im_lf, (sz, sz), interpolation=cv2.INTER_AREA)

            self.digit_render_right.update_object_pose(nut_poses[i])
            self.digit_render_right.update_pose_given_pose(right_finger_poses[i])
            im_rf, hm_rf, mask_rf = self.digit_render_right.render()
            im_rf = self._subtract_bg(im_rf, bg_right)
            im_rf = cv2.resize(im_rf, (sz, sz), interpolation=cv2.INTER_AREA)

            lf[i] = im_lf
            rf[i] = im_rf

        queue.put((lf, rf, offset))

    def refresh_digit_sensors(self):
        if self.t_digits == 0:
            sz = self.cfg["ncf"]["digits_size"]

            left_digit_poses = torch.cat(
                (self.left_fingertip_pos, self.left_fingertip_quat), dim=-1
            )
            right_digit_poses = torch.cat(
                (self.right_fingertip_pos, self.right_fingertip_quat), dim=-1
            )
            nut_poses = torch.cat((self.nut_pos, self.nut_quat), dim=-1)

            tf = np.eye(4)
            tf[0:3, 0:3] = euler_angles_to_matrix(
                euler_angles=torch.tensor([[0, 90, 0]]), convention="XYZ"
            ).numpy()

            left_finger_poses = xyzquat_to_tf_numpy(left_digit_poses.cpu().numpy())
            left_finger_poses = left_finger_poses @ np.array(
                np.linalg.inv(
                    [
                        [0.0, -1.0, 0.0, 0.0],
                        [0.0, 0.0, 1.0, 0.0],
                        [-1.0, 0.0, 0.0, 0.0],
                        [0.0, 0.0, 0.0, 1.0],
                    ]
                )
            )
            left_finger_poses = left_finger_poses @ tf

            right_finger_poses = xyzquat_to_tf_numpy(right_digit_poses.cpu().numpy())
            right_finger_poses = right_finger_poses @ np.array(
                np.linalg.inv(
                    [
                        [0.0, -1.0, 0.0, 0.0],
                        [0.0, 0.0, 1.0, 0.0],
                        [-1.0, 0.0, 0.0, 0.0],
                        [0.0, 0.0, 0.0, 1.0],
                    ]
                )
            )
            right_finger_poses = right_finger_poses @ tf

            nut_poses = xyzquat_to_tf_numpy(nut_poses.cpu().numpy())

            # if self.num_envs > 0:
            groups = 8 if self.num_envs > 8 else 1
            offset = np.ceil(self.num_envs / groups).astype(int)
            processes = []
            results = []
            with Manager() as manager:
                for g in range(groups):
                    env_start = offset * g
                    env_end = np.clip(offset * (g + 1), 0, self.num_envs)
                    nut_poses_g = nut_poses[env_start:env_end]
                    left_finger_poses_g = left_finger_poses[env_start:env_end]
                    right_finger_poses_g = right_finger_poses[env_start:env_end]
                    queue = manager.Queue()
                    process = Process(
                        target=self.render_digit,
                        args=(
                            env_start,
                            nut_poses_g,
                            left_finger_poses_g,
                            right_finger_poses_g,
                            queue,
                        ),
                    )
                    process.start()
                    processes.append(process)
                    results.append(queue)

                for process in processes:
                    process.join()

                for queue in results:
                    left, right, offset = queue.get()
                    self.digit_left_buf[
                        offset : offset + len(left), 0, :
                    ] = torch.tensor(left)
                    self.digit_right_buf[
                        offset : offset + len(right), 0, :
                    ] = torch.tensor(right)

        self.t_digits = (
            self.t_digits + 1 if (self.t_digits < self.subsample_digits - 1) else 0
        )

    def _subtract_bg(self, img1, img2, offset=0.5):
        img1 = np.int32(img1)
        img2 = np.int32(img2)
        diff = img1 - img2
        diff = diff / 255.0 + offset
        return diff

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_env_nut_bolt.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""Factory: abstract base class for base class.

Inherits ABC class. Inherited by base class. Defines template for base class.
"""

from abc import ABC, abstractmethod


class FactoryABCBase(ABC):

    @abstractmethod
    def __init__(self):
        """Initialize instance variables. Initialize VecTask superclass."""
        pass

    @abstractmethod
    def _get_base_yaml_params(self):
        """Initialize instance variables from YAML files."""
        pass

    @abstractmethod
    def create_sim(self):
        """Set sim and PhysX params. Create sim object, ground plane, and envs."""
        pass

    @abstractmethod
    def _create_ground_plane(self):
        """Set ground plane params. Add plane."""
        pass

    @abstractmethod
    def import_franka_assets(self):
        """Set Franka and table asset options. Import assets."""
        pass

    @abstractmethod
    def acquire_base_tensors(self):
        """Acquire and wrap tensors. Create views."""
        pass

    @abstractmethod
    def refresh_base_tensors(self):
        """Refresh tensors."""
        # NOTE: Tensor refresh functions should be called once per step, before setters.
        pass

    @abstractmethod
    def parse_controller_spec(self):
        """Parse controller specification into lower-level controller configuration."""
        pass

    @abstractmethod
    def generate_ctrl_signals(self):
        """Get Jacobian. Set Franka DOF position targets or DOF torques."""
        pass

    @abstractmethod
    def enable_gravity(self):
        """Enable gravity."""
        pass

    @abstractmethod
    def disable_gravity(self):
        """Disable gravity."""
        pass

    @abstractmethod
    def export_scene(self):
        """Export scene to USD."""
        pass

    @abstractmethod
    def extract_poses(self):
        """Extract poses of all bodies."""
        pass

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_class_base.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: abstract base class for environment classes.

Inherits ABC class. Inherited by environment classes. Defines template for environment classes.
"""

from abc import ABC, abstractmethod


class FactoryABCEnv(ABC):

    @abstractmethod
    def __init__(self):
        """Initialize instance variables. Initialize base superclass. Acquire tensors."""
        pass

    @abstractmethod
    def _get_env_yaml_params(self):
        """Initialize instance variables from YAML files."""
        pass

    @abstractmethod
    def create_envs(self):
        """Set env options. Import assets. Create actors."""
        pass

    @abstractmethod
    def _import_env_assets(self):
        """Set asset options. Import assets."""
        pass

    @abstractmethod
    def _create_actors(self):
        """Set initial actor poses. Create actors. Set shape and DOF properties."""
        pass

    @abstractmethod
    def _acquire_env_tensors(self):
        """Acquire and wrap tensors. Create views."""
        pass

    @abstractmethod
    def refresh_env_tensors(self):
        """Refresh tensors."""
        # NOTE: Tensor refresh functions should be called once per step, before setters.
        pass

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_class_env.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: abstract base class for task classes.

Inherits ABC class. Inherited by task classes. Defines template for task classes.
"""

from abc import ABC, abstractmethod


class FactoryABCTask(ABC):

    @abstractmethod
    def __init__(self):
        """Initialize instance variables. Initialize environment superclass."""
        pass

    @abstractmethod
    def _get_task_yaml_params(self):
        """Initialize instance variables from YAML files."""
        pass

    @abstractmethod
    def _acquire_task_tensors(self):
        """Acquire tensors."""
        pass

    @abstractmethod
    def _refresh_task_tensors(self):
        """Refresh tensors."""
        pass

    @abstractmethod
    def pre_physics_step(self):
        """Reset environments. Apply actions from policy as controller targets. Simulation step called after this method."""
        pass

    @abstractmethod
    def post_physics_step(self):
        """Step buffers. Refresh tensors. Compute observations and reward."""
        pass

    @abstractmethod
    def compute_observations(self):
        """Compute observations."""
        pass

    @abstractmethod
    def compute_reward(self):
        """Detect successes and failures. Update reward and reset buffers."""
        pass

    @abstractmethod
    def _update_rew_buf(self):
        """Compute reward at current timestep."""
        pass

    @abstractmethod
    def _update_reset_buf(self):
        """Assign environments for reset if successful or failed."""
        pass

    @abstractmethod
    def reset_idx(self):
        """Reset specified environments."""
        pass

    @abstractmethod
    def _reset_franka(self):
        """Reset DOF states and DOF targets of Franka."""
        pass

    @abstractmethod
    def _reset_object(self):
        """Reset root state of object."""
        pass

    @abstractmethod
    def _reset_buffers(self):
        """Reset buffers."""
        pass

    @abstractmethod
    def _set_viewer_params(self):
        """Set viewer parameters."""
        pass

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_class_task.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: schema for base class configuration.

Used by Hydra. Defines template for base class YAML file.
"""

from dataclasses import dataclass


@dataclass
class Mode:
    export_scene: bool  # export scene to USD
    export_states: bool  # export states to NPY


@dataclass
class PhysX:
    solver_type: int  # default = 1 (Temporal Gauss-Seidel)
    num_threads: int
    num_subscenes: int
    use_gpu: bool
    num_position_iterations: int  # number of position iterations for solver (default = 4)
    num_velocity_iterations: int  # number of velocity iterations for solver (default = 1)
    contact_offset: float  # default = 0.02
    rest_offset: float  # default = 0.001
    bounce_threshold_velocity: float  # default = 0.01
    max_depenetration_velocity: float  # default = 100.0
    friction_offset_threshold: float  # default = 0.04
    friction_correlation_distance: float  # default = 0.025
    max_gpu_contact_pairs: int  # default = 1024 * 1024
    default_buffer_size_multiplier: float
    contact_collection: int  # 0: CC_NEVER (do not collect contact info), 1: CC_LAST_SUBSTEP (collect contact info on last substep), 2: CC_ALL_SUBSTEPS (collect contact info at all substeps)


@dataclass
class Sim:
    dt: float  # timestep size (default = 1.0 / 60.0)
    num_substeps: int  # number of substeps (default = 2)
    up_axis: str
    use_gpu_pipeline: bool
    gravity: list  # gravitational acceleration vector
    add_damping: bool  # add damping to stabilize gripper-object interactions
    physx: PhysX


@dataclass
class Env:
    env_spacing: float  # lateral offset between envs
    franka_depth: float  # depth offset of Franka base relative to env origin
    table_height: float  # height of table
    franka_friction: float  # coefficient of friction associated with Franka
    table_friction: float  # coefficient of friction associated with table


@dataclass
class FactorySchemaConfigBase:
    mode: Mode
    sim: Sim
    env: Env

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_config_base.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: schema for environment class configurations.

Used by Hydra. Defines template for environment class YAML files.
"""

from dataclasses import dataclass


@dataclass
class Sim:
    disable_franka_collisions: bool  # disable collisions between Franka and objects


@dataclass
class Env:
    env_name: str  # name of scene


@dataclass
class FactorySchemaConfigEnv:
    sim: Sim
    env: Env

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_config_env.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Factory: schema for task class configurations.

Used by Hydra. Defines template for task class YAML files. Not enforced.
"""

from __future__ import annotations
from dataclasses import dataclass


@dataclass
class Sim:
    use_gpu_pipeline: bool  # use GPU pipeline
    up_axis: str  # up-down axis {x, y, z}
    dt: float  # timestep size
    gravity: list[float]  # gravity vector

    disable_gravity: bool  # disable gravity for all actors


@dataclass
class Env:
    numObservations: int  # number of observations per env; camel case required by VecTask
    numActions: int  # number of actions per env; camel case required by VecTask
    numEnvs: int  # number of envs; camel case required by VecTask


@dataclass
class Randomize:
    franka_arm_initial_dof_pos: list[float]  # initial Franka arm DOF position (7)


@dataclass
class RL:
    pos_action_scale: list[float]  # scale on pos displacement targets (3), to convert [-1, 1] to +- x m
    rot_action_scale: list[float]  # scale on rot displacement targets (3), to convert [-1, 1] to +- x rad
    force_action_scale: list[float]  # scale on force targets (3), to convert [-1, 1] to +- x N
    torque_action_scale: list[float]  # scale on torque targets (3), to convert [-1, 1] to +- x Nm

    clamp_rot: bool  # clamp small values of rotation actions to zero
    clamp_rot_thresh: float  # smallest acceptable value

    max_episode_length: int  # max number of timesteps in each episode


@dataclass
class All:
    jacobian_type: str  # map between joint space and task space via geometric or analytic Jacobian {geometric, analytic}
    gripper_prop_gains: list[float]  # proportional gains on left and right Franka gripper finger DOF position (2)
    gripper_deriv_gains: list[float]  # derivative gains on left and right Franka gripper finger DOF position (2)


@dataclass
class GymDefault:
    joint_prop_gains: list[int]  # proportional gains on Franka arm DOF position (7)
    joint_deriv_gains: list[int]  # derivative gains on Franka arm DOF position (7)


@dataclass
class JointSpaceIK:
    ik_method: str  # use Jacobian pseudoinverse, Jacobian transpose, damped least squares or adaptive SVD {pinv, trans, dls, svd}
    joint_prop_gains: list[int]
    joint_deriv_gains: list[int]


@dataclass
class JointSpaceID:
    ik_method: str
    joint_prop_gains: list[int]
    joint_deriv_gains: list[int]


@dataclass
class TaskSpaceImpedance:
    motion_ctrl_axes: list[bool]  # axes for which to enable motion control {0, 1} (6)
    task_prop_gains: list[float]  # proportional gains on Franka fingertip pose (6)
    task_deriv_gains: list[float]  # derivative gains on Franka fingertip pose (6)


@dataclass
class OperationalSpaceMotion:
    motion_ctrl_axes: list[bool]
    task_prop_gains: list[float]
    task_deriv_gains: list[float]


@dataclass
class OpenLoopForce:
    force_ctrl_axes: list[bool]  # axes for which to enable force control {0, 1} (6)


@dataclass
class ClosedLoopForce:
    force_ctrl_axes: list[bool]
    wrench_prop_gains: list[float]  # proportional gains on Franka finger force (6)


@dataclass
class HybridForceMotion:
    motion_ctrl_axes: list[bool]
    task_prop_gains: list[float]
    task_deriv_gains: list[float]
    force_ctrl_axes: list[bool]
    wrench_prop_gains: list[float]


@dataclass
class Ctrl:
    ctrl_type: str # {gym_default,
                   #  joint_space_ik,
                   #  joint_space_id,
                   #  task_space_impedance,
                   #  operational_space_motion,
                   #  open_loop_force,
                   #  closed_loop_force,
                   #  hybrid_force_motion}
    gym_default: GymDefault
    joint_space_ik: JointSpaceIK
    joint_space_id: JointSpaceID
    task_space_impedance: TaskSpaceImpedance
    operational_space_motion: OperationalSpaceMotion
    open_loop_force: OpenLoopForce
    closed_loop_force: ClosedLoopForce
    hybrid_force_motion: HybridForceMotion


@dataclass
class FactorySchemaConfigTask:
    name: str
    physics_engine: str
    sim: Sim
    env: Env
    rl: RL
    ctrl: Ctrl

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_schema_config_task.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Class for bowl-in-dishrack task.

Inherits nut-bolt environment class and abstract task class (not enforced).
"""

import os
import math
import torch
import hydra

from isaacgym import gymapi, gymtorch, gymutil
from isaacgymenvs.utils import torch_jit_utils as torch_utils
from isaacgymenvs.tasks.ncf_tacto.factory_task_mug_cupholder import (
    NCFTaskCupholder,
)
from isaacgymenvs.tasks.ncf_tacto.factory_schema_config_env import (
    FactorySchemaConfigEnv,
)
from isaacgymenvs.tasks.ncf_tacto.factory_schema_config_base import (
    FactorySchemaConfigBase,
)
import isaacgymenvs.tasks.ncf_tacto.factory_control as fc
from isaacgymenvs.utils import torch_jit_utils

import trimesh
import numpy as np
from scipy.spatial.transform import Rotation as R
import open3d as o3d


class ExtrinsicContact:
    def __init__(
        self,
        num_envs,
        mesh_obj,
        mesh_dishrack,
        obj_scale,
        dishrack_scale,
        dishrack_pos,
        pointcloud_obj,
        world_ref_pointcloud,
    ) -> None:
        self.num_envs = num_envs
        self.object_trimesh = trimesh.load(mesh_obj)
        self.object_trimesh = self.object_trimesh.apply_scale(obj_scale)
        T = np.eye(4)
        T[0:3, 0:3] = R.from_euler("xyz", [0, 0, 0], degrees=True).as_matrix()
        self.object_trimesh = self.object_trimesh.apply_transform(T)

        self.dishrack_trimesh = trimesh.load(mesh_dishrack)
        # self.dishrack_trimesh = self.dishrack_trimesh.apply_scale(dishrack_scale)
        T = np.eye(4)
        T[0:3, -1] = dishrack_pos
        T[0:3, 0:3] = R.from_euler("xyz", [0.0, 0.0, 3.14], degrees=False).as_matrix()
        self.dishrack_trimesh.apply_transform(T)

        self.dishrack = o3d.t.geometry.RaycastingScene()
        self.dishrack.add_triangles(
            o3d.t.geometry.TriangleMesh.from_legacy(self.dishrack_trimesh.as_open3d)
        )

        # pc = np.load(pointcloud_obj)
        # pc = trimesh.points.PointCloud(pc)
        # T = np.eye(4)
        # T[0:3, 0:3] = R.from_euler("xyz", [0.0, 0.0, -90.0], degrees=True).as_matrix()
        # pc = pc.apply_transform(T)
        # self.pointcloud_obj = pc.vertices
        self.pointcloud_obj = np.load(pointcloud_obj)

        self.n_points = self.pointcloud_obj.shape[0]

        self.gt_extrinsic_contact = torch.zeros((num_envs, self.n_points))
        self.world_ref_pointcloud = world_ref_pointcloud

    def _xyzquat_to_tf_numpy(self, position_quat: np.ndarray) -> np.ndarray:
        """
        convert [x, y, z, qx, qy, qz, qw] to 4 x 4 transformation matrices
        """
        # try:
        position_quat = np.atleast_2d(position_quat)  # (N, 7)
        N = position_quat.shape[0]
        T = np.zeros((N, 4, 4))
        T[:, 0:3, 0:3] = R.from_quat(position_quat[:, 3:]).as_matrix()
        T[:, :3, 3] = position_quat[:, :3]
        T[:, 3, 3] = 1
        # except ValueError:
        #     print("Zero quat error!")
        return T.squeeze()

    def reset_extrinsic_contact(self, env_ids):
        self.gt_extrinsic_contact[env_ids] = torch.zeros((len(env_ids), self.n_points))
        self.step = 0

    def get_extrinsic_contact(self, obj_pos, obj_quat):
        object_poses = torch.cat((obj_pos, obj_quat), dim=1)
        object_poses = self._xyzquat_to_tf_numpy(object_poses.cpu().numpy())

        coords = np.zeros((self.num_envs, self.n_points, 3))
        coords_ref = np.zeros((self.num_envs, self.n_points, 3))
        for i in range(self.num_envs):
            coords_ref[i] = self.pointcloud_obj.copy()
            object_pc_i = trimesh.points.PointCloud(self.pointcloud_obj.copy())
            object_pc_i.apply_transform(object_poses[i])
            coords[i] = np.array(object_pc_i.vertices)

        d = self.dishrack.compute_distance(
            o3d.core.Tensor.from_numpy(coords.astype(np.float32))
        ).numpy()

        c = 0.008
        d = d.flatten()
        idx_2 = np.where(d > c)[0]
        d[idx_2] = c
        d = np.clip(d, 0.0, c)

        d = 1.0 - d / c
        d = np.clip(d, 0.0, 1.0)
        d[d > 0.1] = 1.0
        d = d.reshape((self.num_envs, self.n_points))

        self.gt_extrinsic_contact = torch.tensor(d, dtype=torch.float32)
        coords = torch.tensor(coords, dtype=torch.float32)
        coords_ref = torch.tensor(coords_ref, dtype=torch.float32)

        pointcloud = coords if self.world_ref_pointcloud else coords_ref

        # ncf_mask = torch.where(self.gt_extrinsic_contact > 0.1, 1.0, 0.0)
        # ncf_mask = ncf_mask.repeat(3, 1, 1).permute(1, 2, 0)
        # point_cloud = coords * ncf_mask
        # point_cloud = coords_ref * ncf_mask

        return self.gt_extrinsic_contact, pointcloud


class NCFTaskDishrack(NCFTaskCupholder):
    def __init__(
        self,
        cfg,
        rl_device,
        sim_device,
        graphics_device_id,
        headless,
        virtual_screen_capture,
        force_render,
    ):
        NCFTaskCupholder.__init__(
            self,
            cfg,
            rl_device,
            sim_device,
            graphics_device_id,
            headless,
            virtual_screen_capture,
            force_render,
        )

        current_file = os.path.abspath(__file__)
        self.root_dir = os.path.abspath(
            os.path.join(current_file, "..", "..", "..", "..", "..")
        )

        self.success_cases = 0
        self.failure_cases = 0
        self.result_buf = torch.zeros(0)
        self.num_stats = 1000

        # (
        #     self.ini_ctrl_target_fingertip_midpoint_pos,
        #     self.ini_ctrl_target_fingertip_midpoint_quat,
        # ) = self.load_initial_ee_poses()

        if self.cfg_task.rl.compute_contact_gt:
            dishrack_pos = [
                self.cfg_task.randomize.bolt_pos_xy_initial[0],
                self.cfg_task.randomize.bolt_pos_xy_initial[1],
                self.cfg_base.env.table_height,
            ]
            # re-initialize, dishrack scale should be 1.0
            self.extrinsic_contact_gt = ExtrinsicContact(
                mesh_obj=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_mesh_object
                ),
                mesh_dishrack=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_mesh_dishrack
                ),
                obj_scale=1.0,
                dishrack_scale=1.0,
                dishrack_pos=dishrack_pos,
                pointcloud_obj=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_pointcloud_object
                ),
                num_envs=self.num_envs,
                world_ref_pointcloud=self.cfg_task.rl.world_ref_pointcloud,
            )

    def _get_base_yaml_params(self):
        super()._get_base_yaml_params()
        self.cfg_base.env.franka_friction = 10
        self.cfg_env.env.nut_bolt_friction = 10
        self.cfg_env.env.nut_bolt_friction = 10

    def _get_env_yaml_params(self):
        """Initialize instance variables from YAML files."""
        cs = hydra.core.config_store.ConfigStore.instance()
        cs.store(name="factory_schema_config_env", node=FactorySchemaConfigEnv)

        config_path = (
            "task/NCFEnvNutBolt.yaml"  # relative to Hydra search path (cfg dir)
        )
        self.cfg_env = hydra.compose(config_name=config_path)
        self.cfg_env = self.cfg_env["task"]  # strip superfluous nesting

        asset_info_path = "../../assets/factory/yaml/factory_asset_info_nut_bolt.yaml"
        self.asset_info_nut_bolt = hydra.compose(config_name=asset_info_path)
        # TODO: fix this hardcode and remove previous hardcode hydra code
        self.cfg_env.env["desired_subassemblies"] = ["bowl_dishrack"]
        # strip superfluous nesting
        self.asset_info_nut_bolt = self.asset_info_nut_bolt[""][""][""][""][""][""]
        self.asset_info_nut_bolt = self.asset_info_nut_bolt["assets"]["factory"]["yaml"]

    def _refresh_task_tensors(self):
        """Refresh tensors."""
        # Compute pos of keypoints on gripper, nut, and bolt in world frame
        # hardcoded dishrack target position
        bolt_tip_pos_local = self.bolt_tip_pos_local.clone()
        bolt_tip_pos_local[:, 0] -= 0.12
        bolt_tip_pos_local[:, 1] += 0.05
        bolt_tip_pos_local[:, 2] += 0.01
        for idx, keypoint_offset in enumerate(self.keypoint_offsets):
            self.keypoints_bolt[:, idx] = torch_jit_utils.tf_combine(
                self.bolt_quat,
                self.bolt_pos,
                self.identity_quat,
                (keypoint_offset + bolt_tip_pos_local),
            )[1]

            self.keypoints_nut[:, idx] = torch_jit_utils.tf_combine(
                self.nut_quat,
                self.nut_pos,
                self.identity_quat,
                self.keypoint_offsets_nut[:, idx, :],
            )[1]

    def reset_idx(self, env_ids):
        """Reset specified environments."""
        self._reset_franka(env_ids)
        self._reset_object(env_ids)

        self.refresh_base_tensors()
        self.refresh_env_tensors()
        self._refresh_task_tensors()

        # for i in range(10):
        #     self.gym.simulate(self.sim)
        #     self.render()

        # self.gym.simulate(self.sim)
        # while True:
        #     # self.gym.simulate(self.sim)
        #     self.render()

        # Close gripper onto nut
        # print('start gg')
        # # self.disable_gravity()  # to prevent nut from falling
        # for _ in range(self.cfg_task.env.num_gripper_close_sim_steps * 2):
        #     if self.device == "cpu":
        #         self.gym.fetch_results(self.sim, True)
        #     self.ctrl_target_dof_pos[env_ids, 7:9] = 0.0
        #     delta_hand_pose = torch.zeros(
        #         (self.num_envs, self.cfg_task.env.numActions), device=self.device
        #     )  # no arm motion
        #     self._apply_actions_as_ctrl_targets(
        #         actions=delta_hand_pose, ctrl_target_gripper_dof_pos=0.0, do_scale=False
        #     )
        #     self.gym.simulate(self.sim)
        #     self.render()
        # print('gg')is_nut_close_to_bolt
        # self.enable_gravity(gravity_mag=abs(self.cfg_base.sim.gravity[2]))

        self._randomize_gripper_pose(env_ids, sim_steps=5)

        self._reset_buffers(env_ids)
        if self.cfg_task.rl.log_ncf_data:
            self._init_ncf_logger_tensors(env_ids)
        if self.cfg_task.rl.compute_contact_gt:
            self.extrinsic_contact_gt.reset_extrinsic_contact(env_ids)

    def _get_center_dist(self):
        """Get center distance."""
        # -0.1, 0.05, 0.06
        goals = torch.tensor(
            [[0.0, 0.05, 0.18], [-0.05, 0.05, 0.18], [0.05, 0.05, 0.18]],
            device=self.device,
        )
        center_dist_0 = torch.norm(self.nut_pos - goals[0], p=2, dim=-1)
        center_dist_1 = torch.norm(self.nut_pos - goals[1], p=2, dim=-1)
        center_dist_2 = torch.norm(self.nut_pos - goals[2], p=2, dim=-1)

        center_dist_goals = torch.vstack(
            (center_dist_0, center_dist_1, center_dist_2)
        ).permute(1, 0)

        center_dist = torch.min(center_dist_goals, dim=-1)[0]

        # center_dist = torch.norm(center_dist, p=2, dim=-1)

        return center_dist

    def _check_nut_close_to_bolt(self):
        """Check if nut is close to bolt."""
        # center_dist = self.nut_pos - torch.tensor(
        #     [0.0, 0.065, 0.18], device=self.device
        # )
        # center_dist = torch.norm(center_dist, p=2, dim=-1)
        center_dist = self._get_center_dist()
        is_nut_close_to_bolt = torch.where(
            center_dist < 0.05,
            torch.ones_like(self.progress_buf),
            torch.zeros_like(self.progress_buf),
        )
        return is_nut_close_to_bolt

    def _update_rew_buf(self):
        """Compute reward at current timestep."""
        self._show_debug_viz()

        keypoint_reward = -self._get_center_dist()
        # keypoint_reward = -torch.exp(keypoint_reward * 15)

        # compute error in orientation
        # -0.0075,  0.0332, -0.6507,  0.7586
        # ref_nut_quat = torch.tensor(
        #     [-0.0075, 0.0332, -0.6507, 0.7586], device=self.device
        # ).repeat(len(self.nut_quat), 1)
        # nut_quat_penalty = torch.norm(self.nut_quat - ref_nut_quat, p=2, dim=-1)
        # nut_quat_penalty[nut_quat_penalty < 0.1] = 0.0
        # is_mug_oriented = nut_quat_penalty < 0.1

        dist_nut_bolt = torch.norm(self.bolt_pos - self.nut_pos, p=2, dim=-1)
        # action_penalty = torch.norm(self.actions, p=2, dim=-1)

        self.rew_buf[:] = keypoint_reward * self.cfg_task.rl.keypoint_reward_scale

        # New change: 0.2 instead of 0.1
        # because the center of the bowl might be already off the center of the gripper
        d = torch.norm(self.finger_midpoint_pos - self.nut_pos, p=2, dim=-1)
        is_not_grasped = d >= 0.1
        self.rew_buf[:] += is_not_grasped.float() * -20
        self.reset_buf[is_not_grasped] = 1

        # new reset condition: avoid bowl laying flat on the dishrack (45 degree)
        # r, p, y = torch_utils.get_euler_xyz(self.nut_quat)
        # flipped = torch.logical_and(p > 0.78, p < 5.50)
        # if flipped.any():
        #     self.reset_buf[:] = 1

        # TODO: confirm the following sentence
        # In this policy, episode length is constant across all envs
        # print(self.rew_buf[0], self.progress_buf[0])
        # is_last_step = self.progress_buf[0] == self.max_episode_length - 1
        # if is_last_step:
        # Check if nut is close enough to bolt
        is_nut_close_to_bolt = torch.logical_and(
            self.progress_buf <= self.max_episode_length - 1,
            self._check_nut_close_to_bolt(),
        )

        success = torch.sum(
            torch.logical_and(
                self.progress_buf == self.max_episode_length - 1,
                self._check_nut_close_to_bolt(),
            )
        )
        self.success_cases += success
        failure = torch.sum(
            torch.logical_and(
                self.reset_buf, self.progress_buf != self.max_episode_length - 1
            )
        )
        failure += torch.sum(
            torch.logical_and(
                self.progress_buf == self.max_episode_length - 1,
                torch.logical_not(self._check_nut_close_to_bolt()),
            )
        )
        self.failure_cases += failure

        cur_result_buf = torch.from_numpy(
            np.array([1]) * success.item() + np.array([0]) * failure.item()
        )
        self.result_buf = torch.cat([self.result_buf, cur_result_buf])
        if self.result_buf.shape[0] >= self.num_stats:
            clip_part = self.result_buf.shape[0] - self.num_stats
            self.result_buf = self.result_buf[clip_part:]

        # is_nut_inserted = is_nut_close_to_bolt * is_mug_oriented
        self.rew_buf[:] += is_nut_close_to_bolt * 10  # self.cfg_task.rl.success_bonus
        self.extras["dish_near_slot"] = torch.mean(is_nut_close_to_bolt.float())
        # self.extras["mug_oriented"] = torch.mean(is_mug_oriented.float())
        self.extras["successes"] = self.success_cases / (
            self.success_cases + self.failure_cases + 1e-5
        )
        self.extras["successes2"] = torch.sum(self.result_buf) / (
            self.result_buf.shape[0] + 1e-5
        )
        # self.extras["successes2"] = torch.mean(is_nut_close_to_bolt.float())
        self.extras["dist_dish_dishrack"] = torch.mean(dist_nut_bolt)
        self.extras["keypoint_reward"] = torch.mean(keypoint_reward.abs())
        # self.extras["action_penalty"] = torch.mean(action_penalty)
        # self.extras["mug_quat_penalty"] = torch.mean(nut_quat_penalty)
        self.extras["steps"] = torch.mean(self.progress_buf.float())
        self.extras["dish_close"] = is_nut_close_to_bolt.float()

        if self.reset_buf.sum() > 0:
            if self.cfg_task.rl.log_ncf_data:
                self._save_ncf_data()

    def _reset_franka(self, env_ids, open_gripper=True):
        """Reset DOF states and DOF targets of Franka."""
        nut_widths_max = torch.ones_like(self.nut_widths_max)  # * 0.02
        nut_widths_max = nut_widths_max[env_ids]
        # new change: hard coded because dish is thin
        nut_widths_max[:] = 0.02
        self.dof_pos[env_ids] = torch.cat(
            (
                torch.tensor(
                    self.cfg_task.randomize.franka_arm_initial_dof_pos,
                    device=self.device,
                ).repeat((len(env_ids), 1)),
                (nut_widths_max * 0.5)
                * 1.1,  # buffer on gripper DOF pos to prevent initial contact
                (nut_widths_max * 0.5) * 1.1,
            ),  # buffer on gripper DOF pos to prevent initial contact
            dim=-1,
        )  # shape = (num_envs, num_dofs)

        self.dof_vel[env_ids] = 0.0  # shape = (num_envs, num_dofs)
        self.ctrl_target_dof_pos[env_ids] = self.dof_pos[env_ids]

        multi_env_ids_int32 = self.franka_actor_ids_sim[env_ids].flatten()
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(multi_env_ids_int32),
            len(multi_env_ids_int32),
        )

    def _reset_object(self, env_ids):
        """Reset root states of nut and bolt."""

        # shape of root_pos = (num_envs, num_actors, 3)
        # shape of root_quat = (num_envs, num_actors, 4)
        # shape of root_linvel = (num_envs, num_actors, 3)
        # shape of root_angvel = (num_envs, num_actors, 3)

        # Randomize root state of nut within gripper
        self.root_pos[env_ids, self.nut_actor_id_env, 0] = 0.0
        self.root_pos[env_ids, self.nut_actor_id_env, 1] = 0.0
        fingertip_midpoint_pos_reset = 0.58781  # self.fingertip_midpoint_pos at reset
        # mug
        nut_base_pos_local = 0.005
        self.root_pos[env_ids, self.nut_actor_id_env, 2] = (
            fingertip_midpoint_pos_reset - nut_base_pos_local
        )

        # self.root_pos[env_ids, self.nut_actor_id_env, 1] = 0.02

        # nut
        nut_noise_pos_in_gripper = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        nut_noise_pos_in_gripper = nut_noise_pos_in_gripper @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.nut_noise_pos_in_gripper, device=self.device
            )
        )
        nut_noise_pos_in_gripper[env_ids, 0] -= 0.12
        nut_noise_pos_in_gripper[env_ids, 1] = 0.052
        # this is for not centered dish
        # nut_noise_pos_in_gripper[env_ids, 2] -= 0.205
        nut_noise_pos_in_gripper[env_ids, 2] -= 0.195
        self.root_pos[env_ids, self.nut_actor_id_env, :] += nut_noise_pos_in_gripper[
            env_ids
        ]

        nut_rot_euler = torch.tensor(
            [0.0, 0.0, -math.pi / 2], device=self.device
        ).repeat(len(env_ids), 1)
        nut_noise_rot_in_gripper = 2 * (
            torch.rand(len(env_ids), dtype=torch.float32, device=self.device) - 0.5
        )  # [-1, 1]
        nut_noise_rot_in_gripper *= self.cfg_task.randomize.nut_noise_rot_in_gripper
        nut_rot_euler[:, 2] += nut_noise_rot_in_gripper
        nut_rot_quat = torch_utils.quat_from_euler_xyz(
            nut_rot_euler[:, 0], nut_rot_euler[:, 1], nut_rot_euler[:, 2]
        )
        self.root_quat[env_ids, self.nut_actor_id_env] = nut_rot_quat

        # Randomize root state of bolt
        bolt_noise_xy = 2 * (
            torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        bolt_noise_xy = bolt_noise_xy @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.bolt_pos_xy_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 0] = (
            self.cfg_task.randomize.bolt_pos_xy_initial[0] + bolt_noise_xy[env_ids, 0]
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 1] = (
            self.cfg_task.randomize.bolt_pos_xy_initial[1] + bolt_noise_xy[env_ids, 1]
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 2] = (
            self.cfg_base.env.table_height + self.cfg_task.randomize.bolt_pos_z_offset
        )
        self.root_quat[env_ids, self.bolt_actor_id_env] = torch.tensor(
            [0.0, 0.0, 0.0, 1.0], dtype=torch.float32, device=self.device
        ).repeat(len(env_ids), 1)

        self.root_linvel[env_ids, self.bolt_actor_id_env] = 0.0
        self.root_angvel[env_ids, self.bolt_actor_id_env] = 0.0

        nut_bolt_actor_ids_sim = torch.cat(
            (self.nut_actor_ids_sim[env_ids], self.bolt_actor_ids_sim[env_ids]), dim=0
        )
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(nut_bolt_actor_ids_sim),
            len(nut_bolt_actor_ids_sim),
        )

    def _show_debug_viz(self):
        self.gym.clear_lines(self.viewer)
        sphere_pose = gymapi.Transform()
        sphere_pose.r = gymapi.Quat(0, 0, 0, 1)
        sphere_geom = gymutil.WireframeSphereGeometry(
            0.01, 8, 8, sphere_pose, color=(1, 1, 0)
        )
        sphere_geom_white = gymutil.WireframeSphereGeometry(
            0.02, 8, 8, sphere_pose, color=(1, 1, 1)
        )
        for i in range(self.num_envs):
            if i != 0:
                continue
            palm_center_transform = gymapi.Transform()
            palm_center_transform.p = gymapi.Vec3(*self.nut_pos[0])
            palm_center_transform.r = gymapi.Quat(0, 0, 0, 1)
            gymutil.draw_lines(
                sphere_geom_white,
                self.gym,
                self.viewer,
                self.env_ptrs[i],
                palm_center_transform,
            )

        for i in range(self.num_envs):
            if i != 0:
                continue
            palm_center_transform = gymapi.Transform()
            # palm_center_transform.p = gymapi.Vec3(*[-0.1, 0.05, 0.06])
            palm_center_transform.p = gymapi.Vec3(*[0.0, 0.05, 0.18])
            palm_center_transform.r = gymapi.Quat(0, 0, 0, 1)
            gymutil.draw_lines(
                sphere_geom,
                self.gym,
                self.viewer,
                self.env_ptrs[i],
                palm_center_transform,
            )

    def get_initial_ee_poses(self):
        self.ctrl_target_fingertip_midpoint_pos = self.fingertip_midpoint_pos.clone()

        fingertip_midpoint_pos_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        fingertip_midpoint_pos_noise = fingertip_midpoint_pos_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_pos_noise, device=self.device
            )
        )
        fingertip_midpoint_pos_noise[:, 2] = 0.0
        self.ctrl_target_fingertip_midpoint_pos += fingertip_midpoint_pos_noise

        # Set target rot
        ctrl_target_fingertip_midpoint_euler = (
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_rot_initial,
                device=self.device,
            )
            .unsqueeze(0)
            .repeat(self.num_envs, 1)
        )

        fingertip_midpoint_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        fingertip_midpoint_rot_noise = fingertip_midpoint_rot_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_rot_noise, device=self.device
            )
        )
        ctrl_target_fingertip_midpoint_euler += fingertip_midpoint_rot_noise
        self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_from_euler_xyz(
            ctrl_target_fingertip_midpoint_euler[:, 0],
            ctrl_target_fingertip_midpoint_euler[:, 1],
            ctrl_target_fingertip_midpoint_euler[:, 2],
        )

        torch.save(
            self.ctrl_target_fingertip_midpoint_pos,
            "ctrl_target_fingertip_midpoint_pos.pt",
        )
        torch.save(
            self.ctrl_target_fingertip_midpoint_quat,
            "ctrl_target_fingertip_midpoint_quat.pt",
        )

    def load_initial_ee_poses(self):
        ctrl_target_fingertip_midpoint_pos = torch.load(
            "ctrl_target_fingertip_midpoint_pos.pt"
        )
        ctrl_target_fingertip_midpoint_quat = torch.load(
            "ctrl_target_fingertip_midpoint_quat.pt"
        )
        return ctrl_target_fingertip_midpoint_pos, ctrl_target_fingertip_midpoint_quat

    def set_initial_ee_pose(self, env_ids, sim_steps=10):
        # Step sim and render
        for _ in range(sim_steps):
            self.refresh_base_tensors()
            self.refresh_env_tensors()
            self._refresh_task_tensors()

            pos_error, axis_angle_error = fc.get_pose_error(
                fingertip_midpoint_pos=self.fingertip_midpoint_pos,
                fingertip_midpoint_quat=self.fingertip_midpoint_quat,
                ctrl_target_fingertip_midpoint_pos=self.ini_ctrl_target_fingertip_midpoint_pos,
                ctrl_target_fingertip_midpoint_quat=self.ini_ctrl_target_fingertip_midpoint_quat,
                jacobian_type=self.cfg_ctrl["jacobian_type"],
                rot_error_type="axis_angle",
            )

            delta_hand_pose = torch.cat((pos_error, axis_angle_error), dim=-1)
            actions = torch.zeros(
                (self.num_envs, self.cfg_task.env.numActions), device=self.device
            )
            actions[:, :6] = delta_hand_pose

            self._apply_actions_as_ctrl_targets(
                actions=actions, ctrl_target_gripper_dof_pos=0.0, do_scale=False
            )

            self.gym.simulate(self.sim)
            self.render()

        if self.refresh_digits_flag:
            self.refresh_digit_sensors()

        self.dof_vel[env_ids, :] = torch.zeros_like(self.dof_vel[env_ids])

        # Set DOF state
        multi_env_ids_int32 = self.franka_actor_ids_sim[env_ids].flatten()
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(multi_env_ids_int32),
            len(multi_env_ids_int32),
        )

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_task_bowl_dishrack.py
carolinahiguera/NCF-Policy
# Copyright (c) 2021-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Class for mug-in-cupholder task.

Inherits nut-bolt environment class and abstract task class (not enforced).
"""

import time
import hydra
import math
import omegaconf
import os
import subprocess
import torch
from datetime import datetime
import numpy as np
import matplotlib.pyplot as plt
import cv2
import trimesh
from scipy.spatial.transform import Rotation as R

from isaacgymenvs.utils.torch_jit_utils import quat_apply, to_torch

from isaacgym import gymapi, gymtorch
from isaacgymenvs.utils import torch_jit_utils as torch_utils
import isaacgymenvs.tasks.ncf_tacto.factory_control as fc
from isaacgymenvs.tasks.ncf_tacto.factory_env_nut_bolt import FactoryEnvNutBolt
from isaacgymenvs.tasks.ncf_tacto.factory_schema_class_task import FactoryABCTask
from isaacgymenvs.tasks.ncf_tacto.factory_schema_config_task import (
    FactorySchemaConfigTask,
)
from isaacgymenvs.utils import torch_jit_utils
import open3d as o3d

from multiprocessing import Process, Queue, Manager

# from ncf.digit_vae.models.vae import VAE

DEBUG = False


def get_rank():
    return int(os.getenv("LOCAL_RANK", "0"))


def vulkan_device_id_from_cuda_device_id(orig: int) -> int:
    """Map a CUDA device index to a Vulkan one.

    Used to populate the value of `graphic_device_id`, which in IsaacGym is a vulkan
    device ID.

    This prevents a common segfault we get when the Vulkan ID, which is by default 0,
    points to a device that isn't present in CUDA_VISIBLE_DEVICES.
    """
    # Get UUID of the torch device.
    # All of the private methods can be dropped once this PR lands:
    #     https://github.com/pytorch/pytorch/pull/99967
    try:
        cuda_uuid = torch.cuda._raw_device_uuid_nvml()[
            torch.cuda._parse_visible_devices()[orig]
        ]  # type: ignore
        assert cuda_uuid.startswith("GPU-")
        cuda_uuid = cuda_uuid[4:]
    except AttributeError:
        print("detect cuda / vulkan relation can only be done for pytorch 2.0")
        return get_rank()

    try:
        vulkaninfo_lines = subprocess.run(
            ["vulkaninfo"],
            # We unset DISPLAY to avoid this error:
            # https://github.com/KhronosGroup/Vulkan-Tools/issues/370
            env={k: v for k, v in os.environ.items() if k != "DISPLAY"},
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            universal_newlines=True,
        ).stdout.split("\n")
    except FileNotFoundError:
        print(
            "vulkaninfo was not found; try `apt install vulkan-tools` or `apt install vulkan-utils`."
        )
        return get_rank()

    vulkan_uuids = [
        s.partition("=")[2].strip()
        for s in vulkaninfo_lines
        if s.strip().startswith("deviceUUID")
    ]
    vulkan_uuids = list(dict(zip(vulkan_uuids, vulkan_uuids)).keys())
    vulkan_uuids = [uuid for uuid in vulkan_uuids if not uuid.startswith("0000")]
    out = vulkan_uuids.index(cuda_uuid)
    print(f"Using graphics_device_id={out}", cuda_uuid)
    return out


class ExtrinsicContact:
    def __init__(
        self,
        num_envs,
        mesh_obj,
        mesh_cupholder,
        obj_scale,
        cupholder_scale,
        cupholder_pos,
        pointcloud_obj,
        world_ref_pointcloud,
    ) -> None:
        self.num_envs = num_envs
        self.object_trimesh = trimesh.load(mesh_obj)
        self.object_trimesh = self.object_trimesh.apply_scale(obj_scale)
        T = np.eye(4)
        T[0:3, 0:3] = R.from_euler("xyz", [0, 0, 90], degrees=True).as_matrix()
        self.object_trimesh = self.object_trimesh.apply_transform(T)

        self.cupholder_trimesh = trimesh.load(mesh_cupholder)
        self.cupholder_trimesh = self.cupholder_trimesh.apply_scale(cupholder_scale)
        T = np.eye(4)
        T[0:3, -1] = cupholder_pos
        self.cupholder_trimesh.apply_transform(T)

        self.cupholder = o3d.t.geometry.RaycastingScene()
        self.cupholder.add_triangles(
            o3d.t.geometry.TriangleMesh.from_legacy(self.cupholder_trimesh.as_open3d)
        )

        self.pointcloud_obj = np.load(pointcloud_obj)
        self.n_points = self.pointcloud_obj.shape[0]

        self.gt_extrinsic_contact = torch.zeros((num_envs, self.n_points))
        self.world_ref_pointcloud = world_ref_pointcloud

    def _xyzquat_to_tf_numpy(self, position_quat: np.ndarray) -> np.ndarray:
        """
        convert [x, y, z, qx, qy, qz, qw] to 4 x 4 transformation matrices
        """
        # try:
        position_quat = np.atleast_2d(position_quat)  # (N, 7)
        N = position_quat.shape[0]
        T = np.zeros((N, 4, 4))
        T[:, 0:3, 0:3] = R.from_quat(position_quat[:, 3:]).as_matrix()
        T[:, :3, 3] = position_quat[:, :3]
        T[:, 3, 3] = 1
        # except ValueError:
        #     print("Zero quat error!")
        return T.squeeze()

    def reset_extrinsic_contact(self, env_ids):
        self.gt_extrinsic_contact[env_ids] = torch.zeros((len(env_ids), self.n_points))
        self.step = 0

    def get_extrinsic_contact(self, obj_pos, obj_quat):
        object_poses = torch.cat((obj_pos, obj_quat), dim=1)
        object_poses = self._xyzquat_to_tf_numpy(object_poses.cpu().numpy())

        coords = np.zeros((self.num_envs, self.n_points, 3))
        coords_ref = np.zeros((self.num_envs, self.n_points, 3))
        for i in range(self.num_envs):
            coords_ref[i] = self.pointcloud_obj.copy()
            object_pc_i = trimesh.points.PointCloud(self.pointcloud_obj.copy())
            object_pc_i.apply_transform(object_poses[i])
            coords[i] = np.array(object_pc_i.vertices)

        d = self.cupholder.compute_distance(
            o3d.core.Tensor.from_numpy(coords.astype(np.float32))
        ).numpy()

        c = 0.008
        d = d.flatten()
        idx_2 = np.where(d > c)[0]
        d[idx_2] = c
        d = np.clip(d, 0.0, c)

        d = 1.0 - d / c
        d = np.clip(d, 0.0, 1.0)
        d[d > 0.1] = 1.0
        d = d.reshape((self.num_envs, self.n_points))

        self.gt_extrinsic_contact = torch.tensor(d, dtype=torch.float32)
        coords = torch.tensor(coords, dtype=torch.float32)
        coords_ref = torch.tensor(coords_ref, dtype=torch.float32)

        pointcloud = coords if self.world_ref_pointcloud else coords_ref

        return self.gt_extrinsic_contact, pointcloud


class NCFTaskCupholder(FactoryEnvNutBolt, FactoryABCTask):
    def __init__(
        self,
        cfg,
        rl_device,
        sim_device,
        graphics_device_id,
        headless,
        virtual_screen_capture,
        force_render,
    ):
        """Initialize instance variables. Initialize environment superclass."""

        graphics_device_id = vulkan_device_id_from_cuda_device_id(graphics_device_id)

        super().__init__(
            cfg,
            rl_device,
            sim_device,
            graphics_device_id,
            headless,
            virtual_screen_capture,
            force_render,
        )

        current_file = os.path.abspath(__file__)
        self.root_dir = os.path.abspath(
            os.path.join(current_file, "..", "..", "..", "..", "..")
        )

        self.cfg = cfg
        self._get_task_yaml_params()
        self._acquire_task_tensors()
        self.parse_controller_spec()

        if self.cfg_task.sim.disable_gravity:
            self.disable_gravity()

        if self.viewer is not None:
            self._set_viewer_params()

        if self.cfg_task.rl.log_ncf_data:
            self._init_ncf_logger_tensors()

        # digit sensors modules
        self.refresh_digits_flag = self.cfg_task.sim.sim_digits

        if self.refresh_digits_flag and DEBUG:
            # create plot of 10x2 axes
            self.fig, self.axs = plt.subplots(5, 2, figsize=(7, 20))
            self.fig.tight_layout()
            self.fig.canvas.draw()
            self.fig.canvas.flush_events()

        if self.cfg_task.rl.compute_contact_gt:
            cupholder_pos = [
                self.cfg_task.randomize.bolt_pos_xy_initial[0],
                self.cfg_task.randomize.bolt_pos_xy_initial[1],
                self.cfg_base.env.table_height,
            ]
            self.extrinsic_contact_gt = ExtrinsicContact(
                mesh_obj=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_mesh_object
                ),
                mesh_cupholder=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_mesh_cupholder
                ),
                obj_scale=1.0,
                cupholder_scale=0.75,
                cupholder_pos=cupholder_pos,
                pointcloud_obj=os.path.join(
                    self.root_dir, self.cfg_task.ncf.path_pointcloud_object
                ),
                num_envs=self.num_envs,
                world_ref_pointcloud=self.cfg_task.rl.world_ref_pointcloud,
            )

    def _get_task_yaml_params(self):
        """Initialize instance variables from YAML files."""

        cs = hydra.core.config_store.ConfigStore.instance()
        cs.store(name="factory_schema_config_task", node=FactorySchemaConfigTask)

        self.cfg_task = omegaconf.OmegaConf.create(self.cfg)
        self.max_episode_length = (
            self.cfg_task.rl.max_episode_length
        )  # required instance var for VecTask

        asset_info_path = "../../assets/factory/yaml/factory_asset_info_nut_bolt.yaml"  # relative to Gym's Hydra search path (cfg dir)
        self.cfg_env.env["desired_subassemblies"] = ["mug_cupholder"]
        self.asset_info_nut_bolt = hydra.compose(config_name=asset_info_path)
        self.asset_info_nut_bolt = self.asset_info_nut_bolt[""][""][""][""][""][""][
            "assets"
        ]["factory"][
            "yaml"
        ]  # strip superfluous nesting

        ppo_path = "train/FactoryTaskNutBoltPlacePPO.yaml"  # relative to Gym's Hydra search path (cfg dir)
        self.cfg_ppo = hydra.compose(config_name=ppo_path)
        self.cfg_ppo = self.cfg_ppo["train"]  # strip superfluous nesting

    def _acquire_task_tensors(self):
        """Acquire tensors."""

        # Nut-bolt tensors
        self.nut_base_pos_local = self.bolt_head_heights * torch.tensor(
            [0.0, 0.0, -1.0], device=self.device
        ).repeat((self.num_envs, 1))
        bolt_heights = self.bolt_head_heights + self.bolt_shank_lengths
        self.bolt_tip_pos_local = bolt_heights * torch.tensor(
            [0.0, 0.0, 1.0], device=self.device
        ).repeat((self.num_envs, 1))

        # Keypoint tensors
        self.keypoint_offsets = (
            self._get_keypoint_offsets(self.cfg_task.rl.num_keypoints)
            * self.cfg_task.rl.keypoint_scale
        )
        self.keypoints_nut = torch.zeros(
            (self.num_envs, self.cfg_task.rl.num_keypoints, 3),
            dtype=torch.float32,
            device=self.device,
        )
        self.keypoints_bolt = torch.zeros_like(self.keypoints_nut, device=self.device)

        self.identity_quat = (
            torch.tensor([0.0, 0.0, 0.0, 1.0], device=self.device)
            .unsqueeze(0)
            .repeat(self.num_envs, 1)
        )

        self.actions = torch.zeros(
            (self.num_envs, self.cfg_task.env.numActions), device=self.device
        )

        # keypoints_offset_nut
        self.keypoint_offsets_nut = torch.zeros(
            (self.num_envs, self.cfg_task.rl.num_keypoints, 3), device=self.device
        )
        # nut_heights
        for env_id in range(self.num_envs):
            h = self.nut_heights[env_id].item()
            self.keypoint_offsets_nut[env_id, :, 2] = torch.linspace(
                0, -h, self.cfg_task.rl.num_keypoints, device=self.device
            ).__reversed__()

    def _refresh_task_tensors(self):
        """Refresh tensors."""

        # Compute pos of keypoints on gripper, nut, and bolt in world frame
        bolt_tip_pos_local = self.bolt_tip_pos_local.clone()
        bolt_tip_pos_local[:, 2] -= 0.010
        for idx, keypoint_offset in enumerate(self.keypoint_offsets):
            self.keypoints_bolt[:, idx] = torch_jit_utils.tf_combine(
                self.bolt_quat,
                self.bolt_pos,
                self.identity_quat,
                (keypoint_offset + bolt_tip_pos_local),
            )[1]

            self.keypoints_nut[:, idx] = torch_jit_utils.tf_combine(
                self.nut_quat,
                self.nut_pos,
                self.identity_quat,
                self.keypoint_offsets_nut[:, idx, :],
            )[1]

    def _init_ncf_logger_tensors(self):
        """Initialize tensors for logging."""

        self.ncf_logger = []
        for _ in range(self.num_envs):
            self.env_logger = {}
            self.env_logger["left_fingertip_pos"] = torch.zeros((1, 3))
            self.env_logger["left_fingertip_quat"] = torch.zeros((1, 4))
            self.env_logger["right_fingertip_pos"] = torch.zeros((1, 3))
            self.env_logger["right_fingertip_quat"] = torch.zeros((1, 4))
            self.env_logger["fingertip_centered_pos"] = torch.zeros((1, 3))
            self.env_logger["fingertip_centered_quat"] = torch.zeros((1, 4))
            self.env_logger["mug_pos"] = torch.zeros((1, 3))
            self.env_logger["mug_quat"] = torch.zeros((1, 4))
            self.env_logger["cupholder_pos"] = torch.zeros((1, 3))
            self.env_logger["cupholder_quat"] = torch.zeros((1, 4))
            self.ncf_logger.append(self.env_logger)

    def pre_physics_step(self, actions):
        """Reset environments. Apply actions from policy. Simulation step called after this method."""

        env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
        if len(env_ids) > 0:
            self.reset_idx(env_ids)

        self.actions = actions.clone().to(
            self.device
        )  # shape = (num_envs, num_actions); values = [-1, 1]
        self._apply_actions_as_ctrl_targets(
            actions=self.actions, ctrl_target_gripper_dof_pos=0.0, do_scale=True
        )

    def post_physics_step(self):
        """Step buffers. Refresh tensors. Compute observations and reward. Reset environments."""

        self.progress_buf[:] += 1

        self.refresh_base_tensors()
        self.refresh_env_tensors()

        if self.refresh_digits_flag:
            self.refresh_digit_sensors()

        self._refresh_task_tensors()
        self.compute_observations()
        self.compute_reward()

        if self.cfg_task.rl.compute_contact_gt:
            (
                self.gt_extrinsic_contact,
                self.pointclouds_t,
            ) = self.extrinsic_contact_gt.get_extrinsic_contact(
                obj_pos=self.nut_pos, obj_quat=self.nut_quat
            )

        if self.cfg_task.rl.log_ncf_data:
            self._log_ncf_data(env_ids=torch.arange(self.num_envs))

        if self.cfg_task.rl.debug_viz:
            self._show_debug_viz()
            while True:
                self.render()

        if self.refresh_digits_flag and DEBUG:
            self.plot_digits()

    # debug viz
    def _show_debug_viz(self):
        self.gym.clear_lines(self.viewer)
        self.gym.refresh_rigid_body_state_tensor(self.sim)

        # Grab relevant states to visualize
        eef_pos = self.fingertip_midpoint_pos
        eef_rot = self.fingertip_midpoint_quat
        lf_pos = self.left_fingertip_pos
        lf_rot = self.left_fingertip_quat
        rf_pos = self.right_fingertip_pos
        rf_rot = self.right_fingertip_quat
        mug_pos = self.nut_pos
        mug_rot = self.nut_quat
        cupholder_pos = self.bolt_pos
        cupholder_rot = self.bolt_quat

        # Plot visualizations
        for i in range(self.num_envs):
            # for pos, rot in zip((eef_pos, lf_pos, rf_pos, mug_pos, cupholder_pos), (eef_rot, lf_rot, rf_rot, mug_rot, cupholder_rot)):
            for pos, rot in zip((lf_pos, mug_pos), (lf_rot, mug_rot)):
                px = (
                    (
                        pos[i]
                        + quat_apply(
                            rot[i], to_torch([1, 0, 0], device=self.device) * 0.2
                        )
                    )
                    .cpu()
                    .numpy()
                )
                py = (
                    (
                        pos[i]
                        + quat_apply(
                            rot[i], to_torch([0, 1, 0], device=self.device) * 0.2
                        )
                    )
                    .cpu()
                    .numpy()
                )
                pz = (
                    (
                        pos[i]
                        + quat_apply(
                            rot[i], to_torch([0, 0, 1], device=self.device) * 0.2
                        )
                    )
                    .cpu()
                    .numpy()
                )

                p0 = pos[i].cpu().numpy()
                self.gym.add_lines(
                    self.viewer,
                    self.env_ptrs[i],
                    1,
                    [p0[0], p0[1], p0[2], px[0], px[1], px[2]],
                    [0.85, 0.1, 0.1],
                )
                self.gym.add_lines(
                    self.viewer,
                    self.env_ptrs[i],
                    1,
                    [p0[0], p0[1], p0[2], py[0], py[1], py[2]],
                    [0.1, 0.85, 0.1],
                )
                self.gym.add_lines(
                    self.viewer,
                    self.env_ptrs[i],
                    1,
                    [p0[0], p0[1], p0[2], pz[0], pz[1], pz[2]],
                    [0.1, 0.1, 0.85],
                )

    def compute_observations(self):
        """Compute observations."""
        if self.obs_buf.sum() == 0:
            # initialize obs_buf
            curr_obs = [
                self.fingertip_midpoint_pos,
                self.fingertip_midpoint_quat,
            ]
            curr_obs = torch.cat(curr_obs, dim=-1)
            self.obs_buf = torch.cat((curr_obs, curr_obs, curr_obs), dim=-1)
        else:
            curr_obs = [
                self.fingertip_midpoint_pos,
                self.fingertip_midpoint_quat,
            ]
            curr_obs = torch.cat(curr_obs, dim=-1)
            prev_obs_buf = self.obs_buf[:, 0:14].clone()
            self.obs_buf = torch.cat((curr_obs, prev_obs_buf), dim=-1)

    def compute_observations_old(self):
        """Compute observations."""

        # add noise to bolt position and orientation
        bolt_pos_noise = (2 * torch.randn_like(self.bolt_pos) - 1) * 0.01  # max 1cm
        bolt_pos_noise[:, 2] = 0  # no noise in z direction
        bolt_quat_noise = (2 * torch.randn_like(self.bolt_quat) - 1) * 0.0

        # Shallow copies of tensors
        if self.cfg_task.rl.reduce_obs:
            obs_tensors = [
                self.fingertip_midpoint_pos,  # (3)
                self.fingertip_midpoint_quat,  # (4)
                self.bolt_pos + bolt_pos_noise,  # (3)
                self.bolt_quat + bolt_quat_noise,  # (4)
            ]
        else:
            obs_tensors = [
                self.fingertip_midpoint_pos,  # (3)
                self.fingertip_midpoint_quat,  # (4)
                self.fingertip_midpoint_linvel,  # (3)
                self.fingertip_midpoint_angvel,  # (3)
                self.nut_pos,  # (3)
                self.nut_quat,  # (4)
                self.bolt_pos,  # (3)
                self.bolt_quat,
            ]  # (4)

        if self.cfg_task.rl.add_obs_bolt_tip_pos:
            obs_tensors += [self.bolt_tip_pos_local]

        self.obs_buf = torch.cat(
            obs_tensors, dim=-1
        )  # shape = (num_envs, num_observations)

        return self.obs_buf

    def compute_reward(self):
        """Update reward and reset buffers."""

        self._update_reset_buf()
        self._update_rew_buf()

    def _update_reset_buf(self):
        """Assign environments for reset if successful or failed."""

        # If max episode length has been reached
        self.reset_buf[:] = torch.where(
            self.progress_buf[:] >= self.cfg_task.rl.max_episode_length - 1,
            torch.ones_like(self.reset_buf),
            self.reset_buf,
        )

    def _update_rew_buf(self):
        """Compute reward at current timestep."""

        if self.cfg_task.rl.use_keypoints_rew:
            keypoint_reward = -self._get_keypoint_dist()
        else:
            keypoint_reward = -self._get_center_dist()

        # compute error in orientation
        ref_nut_quat = torch.tensor(
            [-0.0075, 0.0332, -0.6507, 0.7586], device=self.device
        ).repeat(len(self.nut_quat), 1)
        nut_quat_penalty = torch.norm(self.nut_quat - ref_nut_quat, p=2, dim=-1)
        nut_quat_penalty[nut_quat_penalty < 0.1] = 0.0
        is_mug_oriented = nut_quat_penalty < 0.1

        dist_nut_bolt = torch.norm(self.bolt_pos - self.nut_pos, p=2, dim=-1)

        action_penalty = torch.norm(self.actions, p=2, dim=-1)

        self.rew_buf[:] = (
            keypoint_reward * self.cfg_task.rl.keypoint_reward_scale
            - action_penalty * self.cfg_task.rl.action_penalty_scale
            - nut_quat_penalty * self.cfg_task.rl.orientation_penalty_scale
            - dist_nut_bolt * self.cfg_task.rl.dist_penalty_scale
        )

        # check is object is grasped and reset if not
        d = torch.norm(self.finger_midpoint_pos - self.nut_pos, p=2, dim=-1)
        is_not_grasped = d >= 0.10
        self.reset_buf[is_not_grasped] = 1

        is_nut_close_to_bolt = self._check_nut_close_to_bolt()
        is_nut_inserted = is_nut_close_to_bolt * is_mug_oriented
        tmp = is_nut_inserted * self.progress_buf
        self.time_complete_task[self.time_complete_task == 0] = tmp[
            self.time_complete_task == 0
        ]

        # In this policy, episode length is constant across all envs
        is_last_step = self.progress_buf[0] == self.max_episode_length - 1
        if is_last_step:
            self.time_complete_task[self.time_complete_task == 0] = self.progress_buf[
                self.time_complete_task == 0
            ]
            # Check if nut is close enough to bolt
            is_nut_close_to_bolt = self._check_nut_close_to_bolt()
            is_nut_inserted = is_nut_close_to_bolt * is_mug_oriented
            self.rew_buf[:] += is_nut_inserted * self.cfg_task.rl.success_bonus
            self.extras["mug_close_to_cupholder"] = torch.mean(
                is_nut_close_to_bolt.float()
            )
            self.extras["mug_oriented"] = torch.mean(is_mug_oriented.float())
            self.extras["successes"] = torch.mean(is_nut_inserted.float())
            self.extras["dist_mug_cupholder"] = torch.mean(dist_nut_bolt)
            self.extras["keypoint_reward"] = torch.mean(keypoint_reward.abs())
            self.extras["action_penalty"] = torch.mean(action_penalty)
            self.extras["mug_quat_penalty"] = torch.mean(nut_quat_penalty)
            self.extras["steps"] = torch.mean(self.progress_buf.float())
            self.extras["mean_time_complete_task"] = torch.mean(
                self.time_complete_task.float()
            )
            a = self.time_complete_task.float() * is_nut_inserted
            self.extras["time_success_task"] = a.sum() / torch.where(a > 0)[0].shape[0]

            if True in is_nut_close_to_bolt:
                print("nut is close to bolt")

            if self.cfg_task.rl.log_ncf_data:
                self._save_ncf_data()

    def _log_ncf_data(self, env_ids):
        for env_id in env_ids:
            self.ncf_logger[env_id]["left_fingertip_pos"] = torch.cat(
                (
                    self.ncf_logger[env_id]["left_fingertip_pos"],
                    self.left_fingertip_pos[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["left_fingertip_quat"] = torch.cat(
                (
                    self.ncf_logger[env_id]["left_fingertip_quat"],
                    self.left_fingertip_quat[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["right_fingertip_pos"] = torch.cat(
                (
                    self.ncf_logger[env_id]["right_fingertip_pos"],
                    self.right_fingertip_pos[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["right_fingertip_quat"] = torch.cat(
                (
                    self.ncf_logger[env_id]["right_fingertip_quat"],
                    self.right_fingertip_quat[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["fingertip_centered_pos"] = torch.cat(
                (
                    self.ncf_logger[env_id]["fingertip_centered_pos"],
                    self.fingertip_midpoint_pos[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["fingertip_centered_quat"] = torch.cat(
                (
                    self.ncf_logger[env_id]["fingertip_centered_quat"],
                    self.fingertip_midpoint_quat[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["mug_pos"] = torch.cat(
                (
                    self.ncf_logger[env_id]["mug_pos"],
                    self.nut_pos[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["mug_quat"] = torch.cat(
                (
                    self.ncf_logger[env_id]["mug_quat"],
                    self.nut_quat[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["cupholder_pos"] = torch.cat(
                (
                    self.ncf_logger[env_id]["cupholder_pos"],
                    self.bolt_pos[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )
            self.ncf_logger[env_id]["cupholder_quat"] = torch.cat(
                (
                    self.ncf_logger[env_id]["cupholder_quat"],
                    self.bolt_quat[[env_id]].unsqueeze(0).cpu(),
                ),
                dim=0,
            )

    def _save_ncf_data(self):
        path_save_dataset = self.cfg_task.rl.path_save_ncf_data
        torch.save(
            self.ncf_logger,
            f"{path_save_dataset}/ncf_data_{datetime.now():%m-%d %H-%M-%S}.pt",
        )

    def reset_idx(self, env_ids):
        """Reset specified environments."""

        self._reset_franka(env_ids)
        # self._reset_object_debug(env_ids)
        self._reset_object(env_ids)

        self.refresh_base_tensors()
        self.refresh_env_tensors()
        self._refresh_task_tensors()
        self.time_complete_task = torch.zeros_like(self.progress_buf)

        # Close gripper onto nut
        self.disable_gravity()  # to prevent nut from falling
        for _ in range(self.cfg_task.env.num_gripper_close_sim_steps):
            self.ctrl_target_dof_pos[env_ids, 7:9] = 0.0
            delta_hand_pose = torch.zeros(
                (self.num_envs, self.cfg_task.env.numActions), device=self.device
            )  # no arm motion
            self._apply_actions_as_ctrl_targets(
                actions=delta_hand_pose, ctrl_target_gripper_dof_pos=0.0, do_scale=False
            )
            self.gym.simulate(self.sim)
            self.render()

        self.enable_gravity(gravity_mag=abs(self.cfg_base.sim.gravity[2]))
        self._move_down_arm(
            env_ids, sim_steps=self.cfg_task.env.num_gripper_move_sim_steps * 2
        )
        self._randomize_gripper_pose(
            env_ids, sim_steps=self.cfg_task.env.num_gripper_move_sim_steps
        )
        self._reset_buffers(env_ids)

        if self.cfg_task.rl.log_ncf_data:
            self._init_ncf_logger_tensors()
        if self.cfg_task.rl.compute_contact_gt:
            self.extrinsic_contact_gt.reset_extrinsic_contact(env_ids)

    def _reset_franka(self, env_ids, open_gripper=True):
        """Reset DOF states and DOF targets of Franka."""
        nut_widths_max = torch.ones_like(self.nut_widths_max)  # * 0.02
        nut_widths_max = nut_widths_max[env_ids]
        self.dof_pos[env_ids] = torch.cat(
            (
                torch.tensor(
                    self.cfg_task.randomize.franka_arm_initial_dof_pos,
                    device=self.device,
                ).repeat((len(env_ids), 1)),
                (nut_widths_max * 0.5)
                * 1.1,  # buffer on gripper DOF pos to prevent initial contact
                (nut_widths_max * 0.5) * 1.1,
            ),  # buffer on gripper DOF pos to prevent initial contact
            dim=-1,
        )  # shape = (num_envs, num_dofs)

        self.dof_vel[env_ids] = 0.0  # shape = (num_envs, num_dofs)
        self.ctrl_target_dof_pos[env_ids] = self.dof_pos[env_ids]

        multi_env_ids_int32 = self.franka_actor_ids_sim[env_ids].flatten()
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(multi_env_ids_int32),
            len(multi_env_ids_int32),
        )

    def _reset_object_debug(self, env_ids):
        """Reset root states of nut and bolt."""

        # shape of root_pos = (num_envs, num_actors, 3)
        # shape of root_quat = (num_envs, num_actors, 4)
        # shape of root_linvel = (num_envs, num_actors, 3)
        # shape of root_angvel = (num_envs, num_actors, 3)

        # Randomize root state of nut within gripper
        self.root_pos[env_ids, self.nut_actor_id_env, 0] = 0.0
        self.root_pos[env_ids, self.nut_actor_id_env, 1] = 0.0
        self.root_pos[env_ids, self.nut_actor_id_env, 2] = 0.06

        nut_rot_euler = torch.tensor(
            [0.0, 0.0, -math.pi / 2], device=self.device
        ).repeat(len(env_ids), 1)
        nut_rot_quat = torch_utils.quat_from_euler_xyz(
            nut_rot_euler[:, 0], nut_rot_euler[:, 1], nut_rot_euler[:, 2]
        )
        self.root_quat[env_ids, self.nut_actor_id_env] = nut_rot_quat

        # Randomize root state of bolt
        self.root_pos[
            env_ids, self.bolt_actor_id_env, 0
        ] = self.cfg_task.randomize.bolt_pos_xy_initial[0]
        self.root_pos[
            env_ids, self.bolt_actor_id_env, 1
        ] = self.cfg_task.randomize.bolt_pos_xy_initial[1]
        self.root_pos[
            env_ids, self.bolt_actor_id_env, 2
        ] = self.cfg_base.env.table_height
        self.root_quat[env_ids, self.bolt_actor_id_env] = torch.tensor(
            [0.0, 0.0, 0.0, 1.0], dtype=torch.float32, device=self.device
        ).repeat(len(env_ids), 1)

        self.root_linvel[env_ids, self.bolt_actor_id_env] = 0.0
        self.root_angvel[env_ids, self.bolt_actor_id_env] = 0.0

        nut_bolt_actor_ids_sim = torch.cat(
            (self.nut_actor_ids_sim[env_ids], self.bolt_actor_ids_sim[env_ids]), dim=0
        )
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(nut_bolt_actor_ids_sim),
            len(nut_bolt_actor_ids_sim),
        )

    def _reset_object(self, env_ids):
        """Reset root states of nut and bolt."""

        # shape of root_pos = (num_envs, num_actors, 3)
        # shape of root_quat = (num_envs, num_actors, 4)
        # shape of root_linvel = (num_envs, num_actors, 3)
        # shape of root_angvel = (num_envs, num_actors, 3)

        # Randomize root state of nut within gripper
        self.root_pos[env_ids, self.nut_actor_id_env, 0] = 0.0
        self.root_pos[env_ids, self.nut_actor_id_env, 1] = 0.0
        fingertip_midpoint_pos_reset = 0.58781  # self.fingertip_midpoint_pos at reset
        # mug
        nut_base_pos_local = 0.005
        self.root_pos[env_ids, self.nut_actor_id_env, 2] = (
            fingertip_midpoint_pos_reset - nut_base_pos_local
        )

        # self.root_pos[env_ids, self.nut_actor_id_env, 1] = 0.02

        # nut
        nut_noise_pos_in_gripper = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        nut_noise_pos_in_gripper = nut_noise_pos_in_gripper @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.nut_noise_pos_in_gripper, device=self.device
            )
        )
        self.root_pos[env_ids, self.nut_actor_id_env, :] += nut_noise_pos_in_gripper[
            env_ids
        ]

        nut_rot_euler = torch.tensor(
            [0.0, 0.0, -math.pi / 2], device=self.device
        ).repeat(len(env_ids), 1)
        nut_noise_rot_in_gripper = 2 * (
            torch.rand(len(env_ids), dtype=torch.float32, device=self.device) - 0.5
        )  # [-1, 1]
        nut_noise_rot_in_gripper *= self.cfg_task.randomize.nut_noise_rot_in_gripper
        nut_rot_euler[:, 2] += nut_noise_rot_in_gripper
        nut_rot_quat = torch_utils.quat_from_euler_xyz(
            nut_rot_euler[:, 0], nut_rot_euler[:, 1], nut_rot_euler[:, 2]
        )
        self.root_quat[env_ids, self.nut_actor_id_env] = nut_rot_quat

        # Randomize root state of bolt
        bolt_noise_xy = 2 * (
            torch.rand((self.num_envs, 2), dtype=torch.float32, device=self.device)
            - 0.5
        )  # [-1, 1]
        bolt_noise_xy = bolt_noise_xy @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.bolt_pos_xy_noise,
                dtype=torch.float32,
                device=self.device,
            )
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 0] = (
            self.cfg_task.randomize.bolt_pos_xy_initial[0] + bolt_noise_xy[env_ids, 0]
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 1] = (
            self.cfg_task.randomize.bolt_pos_xy_initial[1] + bolt_noise_xy[env_ids, 1]
        )
        self.root_pos[env_ids, self.bolt_actor_id_env, 2] = (
            self.cfg_base.env.table_height + self.cfg_task.randomize.bolt_pos_z_offset
        )
        self.root_quat[env_ids, self.bolt_actor_id_env] = torch.tensor(
            [0.0, 0.0, 0.0, 1.0], dtype=torch.float32, device=self.device
        ).repeat(len(env_ids), 1)

        self.root_linvel[env_ids, self.bolt_actor_id_env] = 0.0
        self.root_angvel[env_ids, self.bolt_actor_id_env] = 0.0

        nut_bolt_actor_ids_sim = torch.cat(
            (self.nut_actor_ids_sim[env_ids], self.bolt_actor_ids_sim[env_ids]), dim=0
        )
        self.gym.set_actor_root_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.root_state),
            gymtorch.unwrap_tensor(nut_bolt_actor_ids_sim),
            len(nut_bolt_actor_ids_sim),
        )

    def _reset_buffers(self, env_ids):
        """Reset buffers."""

        self.reset_buf[env_ids] = 0
        self.progress_buf[env_ids] = 0

    def _set_viewer_params(self):
        """Set viewer parameters."""

        cam_pos = gymapi.Vec3(-0.3, -0.3, 0.3)
        cam_target = gymapi.Vec3(0.0, 0.5, 0.01)
        self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

    def _apply_actions_as_ctrl_targets(
        self, actions, ctrl_target_gripper_dof_pos, do_scale
    ):
        """Apply actions from policy as position/rotation targets."""

        # Interpret actions as target pos displacements and set pos target
        pos_actions = actions[:, 0:3]
        if do_scale:
            pos_actions = pos_actions @ torch.diag(
                torch.tensor(self.cfg_task.rl.pos_action_scale, device=self.device)
            )
        self.ctrl_target_fingertip_midpoint_pos = self.fingertip_midpoint_pos + (
            pos_actions
        )

        # Interpret actions as target rot (axis-angle) displacements
        rot_actions = actions[:, 3:6]  # * 0.5
        if do_scale:
            rot_actions = rot_actions @ torch.diag(
                torch.tensor(self.cfg_task.rl.rot_action_scale, device=self.device)
            )

        # Convert to quat and set rot target
        angle = torch.norm(rot_actions, p=2, dim=-1)
        axis = rot_actions / angle.unsqueeze(-1)
        rot_actions_quat = torch_utils.quat_from_angle_axis(angle, axis)
        if self.cfg_task.rl.clamp_rot:
            rot_actions_quat = torch.where(
                angle.unsqueeze(-1).repeat(1, 4) > self.cfg_task.rl.clamp_rot_thresh,
                rot_actions_quat,
                torch.tensor([0.0, 0.0, 0.0, 1.0], device=self.device).repeat(
                    self.num_envs, 1
                ),
            )
        self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_mul(
            rot_actions_quat, self.fingertip_midpoint_quat
        )

        if self.cfg_ctrl["do_force_ctrl"]:
            # Interpret actions as target forces and target torques
            force_actions = actions[:, 6:9]
            if do_scale:
                force_actions = force_actions @ torch.diag(
                    torch.tensor(
                        self.cfg_task.rl.force_action_scale, device=self.device
                    )
                )

            torque_actions = actions[:, 9:12]
            if do_scale:
                torque_actions = torque_actions @ torch.diag(
                    torch.tensor(
                        self.cfg_task.rl.torque_action_scale, device=self.device
                    )
                )

            self.ctrl_target_fingertip_contact_wrench = torch.cat(
                (force_actions, torque_actions), dim=-1
            )

        self.ctrl_target_gripper_dof_pos = ctrl_target_gripper_dof_pos

        self.generate_ctrl_signals()

    def _open_gripper(self, sim_steps=20):
        """Fully open gripper using controller. Called outside RL loop (i.e., after last step of episode)."""

        self._move_gripper_to_dof_pos(gripper_dof_pos=0.1, sim_steps=sim_steps)

    def _move_gripper_to_dof_pos(self, gripper_dof_pos, sim_steps=20):
        """Move gripper fingers to specified DOF position using controller."""

        delta_hand_pose = torch.zeros(
            (self.num_envs, self.cfg_task.env.numActions), device=self.device
        )  # no arm motion
        self._apply_actions_as_ctrl_targets(
            delta_hand_pose, gripper_dof_pos, do_scale=False
        )

        # Step sim
        for _ in range(sim_steps):
            self.render()
            self.gym.simulate(self.sim)

    def _lift_gripper(self, gripper_dof_pos=0.0, lift_distance=0.3, sim_steps=20):
        """Lift gripper by specified distance. Called outside RL loop (i.e., after last step of episode)."""

        delta_hand_pose = torch.zeros([self.num_envs, 6], device=self.device)
        delta_hand_pose[:, 2] = lift_distance  # lift along z

        # Step sim
        for _ in range(sim_steps):
            self._apply_actions_as_ctrl_targets(
                delta_hand_pose, gripper_dof_pos, do_scale=False
            )
            self.render()
            self.gym.simulate(self.sim)

    def _get_keypoint_offsets(self, num_keypoints):
        """Get uniformly-spaced keypoints along a line of unit length, centered at 0."""

        keypoint_offsets = torch.zeros((num_keypoints, 3), device=self.device)
        keypoint_offsets[:, -1] = torch.linspace(
            0.0, 1.0, num_keypoints, device=self.device
        )

        return keypoint_offsets

    def _get_keypoint_dist(self):
        """Get keypoint distances."""
        keypoint_dist = torch.mean(
            torch.norm(self.keypoints_bolt - self.keypoints_nut, p=2, dim=-1), dim=-1
        )
        return keypoint_dist

    def _get_center_dist(self):
        """Get center distance."""

        nut_pos = self.nut_pos - torch.tensor([0.0, 0.0, 0.04], device=self.device)
        center_dist = torch.norm(self.bolt_pos - nut_pos, p=2, dim=-1)

        return center_dist

    def _check_nut_close_to_bolt(self):
        """Check if nut is close to bolt."""
        keypoint_dist = torch.norm(
            self.keypoints_bolt - self.keypoints_nut, p=2, dim=-1
        )

        is_nut_close_to_bolt = torch.where(
            torch.mean(keypoint_dist, dim=-1) < self.cfg_task.rl.close_error_thresh,
            torch.ones_like(self.progress_buf),
            torch.zeros_like(self.progress_buf),
        )

        return is_nut_close_to_bolt

    def _check_nut_away_from_bolt(self):
        """Check if nut is away from bolt."""

        keypoint_dist = torch.norm(
            self.keypoints_bolt - self.keypoints_nut, p=2, dim=-1
        )

        is_nut_away_from_bolt = torch.where(
            torch.mean(keypoint_dist, dim=-1) >= self.cfg_task.rl.away_error_thresh,
            torch.ones_like(self.progress_buf),
            torch.zeros_like(self.progress_buf),
        )

        return is_nut_away_from_bolt

    def _move_down_arm(self, env_ids, sim_steps=20):
        self.ctrl_target_fingertip_midpoint_pos = torch.tensor(
            self.cfg_task.randomize.fingertip_midpoint_pos_initial,
            device=self.device,
        ) + torch.tensor([0.0, 0.00, -0.05], device=self.device)

        self.ctrl_target_fingertip_midpoint_pos = (
            self.ctrl_target_fingertip_midpoint_pos.unsqueeze(0).repeat(
                self.num_envs, 1
            )
        )

        ctrl_target_fingertip_midpoint_euler = (
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_rot_initial,
                device=self.device,
            )
            .unsqueeze(0)
            .repeat(self.num_envs, 1)
        )
        self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_from_euler_xyz(
            ctrl_target_fingertip_midpoint_euler[:, 0],
            ctrl_target_fingertip_midpoint_euler[:, 1],
            ctrl_target_fingertip_midpoint_euler[:, 2],
        )

        for _ in range(sim_steps):
            self.refresh_base_tensors()
            self.refresh_env_tensors()
            self._refresh_task_tensors()

            pos_error, axis_angle_error = fc.get_pose_error(
                fingertip_midpoint_pos=self.fingertip_midpoint_pos,
                fingertip_midpoint_quat=self.fingertip_midpoint_quat,
                ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos,
                ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat,
                jacobian_type=self.cfg_ctrl["jacobian_type"],
                rot_error_type="axis_angle",
            )

            delta_hand_pose = torch.cat((pos_error, axis_angle_error), dim=-1)
            actions = torch.zeros(
                (self.num_envs, self.cfg_task.env.numActions), device=self.device
            )
            actions[:, :6] = delta_hand_pose

            self._apply_actions_as_ctrl_targets(
                actions=actions, ctrl_target_gripper_dof_pos=0.0, do_scale=False
            )

            self.gym.simulate(self.sim)
            self.render()

        if self.refresh_digits_flag:
            self.refresh_digit_sensors()

        self.dof_vel[env_ids, :] = torch.zeros_like(self.dof_vel[env_ids])

        # Set DOF state
        multi_env_ids_int32 = self.franka_actor_ids_sim[env_ids].flatten()
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(multi_env_ids_int32),
            len(multi_env_ids_int32),
        )

    def _randomize_gripper_pose(self, env_ids, sim_steps):
        """Move gripper to random pose."""

        # Set target pos above table
        self.ctrl_target_fingertip_midpoint_pos = self.fingertip_midpoint_pos.clone()

        self.ctrl_target_fingertip_midpoint_pos = torch.tensor(
            [0.0, 0.0, self.cfg_base.env.table_height], device=self.device
        ) + torch.tensor(
            self.cfg_task.randomize.fingertip_midpoint_pos_initial, device=self.device
        )
        self.ctrl_target_fingertip_midpoint_pos = (
            self.ctrl_target_fingertip_midpoint_pos.unsqueeze(0).repeat(
                self.num_envs, 1
            )
        )

        fingertip_midpoint_pos_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )
        fingertip_midpoint_pos_noise = fingertip_midpoint_pos_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_pos_noise, device=self.device
            )
        )
        fingertip_midpoint_pos_noise[:, 2] = 0.0
        self.ctrl_target_fingertip_midpoint_pos += fingertip_midpoint_pos_noise

        # Set target rot
        ctrl_target_fingertip_midpoint_euler = (
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_rot_initial,
                device=self.device,
            )
            .unsqueeze(0)
            .repeat(self.num_envs, 1)
        )

        fingertip_midpoint_rot_noise = 2 * (
            torch.rand((self.num_envs, 3), dtype=torch.float32, device=self.device)
            - 0.5
        )
        fingertip_midpoint_rot_noise = fingertip_midpoint_rot_noise @ torch.diag(
            torch.tensor(
                self.cfg_task.randomize.fingertip_midpoint_rot_noise, device=self.device
            )
        )
        ctrl_target_fingertip_midpoint_euler += fingertip_midpoint_rot_noise
        self.ctrl_target_fingertip_midpoint_quat = torch_utils.quat_from_euler_xyz(
            ctrl_target_fingertip_midpoint_euler[:, 0],
            ctrl_target_fingertip_midpoint_euler[:, 1],
            ctrl_target_fingertip_midpoint_euler[:, 2],
        )

        # Step sim and render
        for _ in range(sim_steps):
            self.refresh_base_tensors()
            self.refresh_env_tensors()
            self._refresh_task_tensors()

            pos_error, axis_angle_error = fc.get_pose_error(
                fingertip_midpoint_pos=self.fingertip_midpoint_pos,
                fingertip_midpoint_quat=self.fingertip_midpoint_quat,
                ctrl_target_fingertip_midpoint_pos=self.ctrl_target_fingertip_midpoint_pos,
                ctrl_target_fingertip_midpoint_quat=self.ctrl_target_fingertip_midpoint_quat,
                jacobian_type=self.cfg_ctrl["jacobian_type"],
                rot_error_type="axis_angle",
            )

            delta_hand_pose = torch.cat((pos_error, axis_angle_error), dim=-1)
            actions = torch.zeros(
                (self.num_envs, self.cfg_task.env.numActions), device=self.device
            )
            actions[:, :6] = delta_hand_pose

            self._apply_actions_as_ctrl_targets(
                actions=actions, ctrl_target_gripper_dof_pos=0.0, do_scale=False
            )

            self.gym.simulate(self.sim)
            self.render()

        if self.refresh_digits_flag:
            self.refresh_digit_sensors()

        self.dof_vel[env_ids, :] = torch.zeros_like(self.dof_vel[env_ids])

        # Set DOF state
        multi_env_ids_int32 = self.franka_actor_ids_sim[env_ids].flatten()
        self.gym.set_dof_state_tensor_indexed(
            self.sim,
            gymtorch.unwrap_tensor(self.dof_state),
            gymtorch.unwrap_tensor(multi_env_ids_int32),
            len(multi_env_ids_int32),
        )

    def plot_digits(self):
        for ax in self.axs.flatten():
            ax.clear()

        for env_id, ax in enumerate(self.axs.flatten()):
            img_left = self.digit_left_buf[env_id][0].cpu().numpy()
            img_right = self.digit_right_buf[env_id][0].cpu().numpy()
            img_left = cv2.resize(img_left, (240, 320), interpolation=cv2.INTER_LINEAR)
            img_right = cv2.resize(
                img_right, (240, 320), interpolation=cv2.INTER_LINEAR
            )
            # img_left = np.random.random(size=(64, 64, 3))
            img = np.concatenate((img_left, img_right), axis=1)
            ax.imshow(img)
            ax.axis("off")
            ax.set_title(f"env_id: {env_id}")

        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        plt.pause(0.001)

NCF_policies/isaacgymenvs/tasks/ncf_tacto/factory_task_mug_cupholder.py
carolinahiguera/NCF-Policy
# Copyright (c) Meta Platforms, Inc. and affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# Extension of Hora policy to tactile sensing with the DIGIT
# Authors: Haozhi Qi, Sudharshan Suresh
# --------------------------------------------------------

"""
    SE(3) pose utilities 
"""

# from isaacgym import gymapi
import numpy as np
import torch
from scipy.spatial.transform import Rotation as R
import copy 

def tf_to_xyzquat_numpy(pose: torch.Tensor) -> torch.Tensor:
    """
    convert 4 x 4 transformation matrices to [x, y, z, qx, qy, qz, qw]
    """
    pose = np.atleast_3d(pose)

    r = R.from_matrix(np.array(pose[:, 0:3, 0:3]))
    q = r.as_quat()  # qx, qy, qz, qw
    t = pose[:, :3, 3]
    xyz_quat = np.concatenate((t, q), axis=1)

    return xyz_quat  # (N, 7)


def xyzquat_to_tf_numpy(position_quat: np.ndarray) -> np.ndarray:
    """
    convert [x, y, z, qx, qy, qz, qw] to 4 x 4 transformation matrices
    """
    # try:
    position_quat = np.atleast_2d(position_quat)  # (N, 7)
    N = position_quat.shape[0]
    T = np.zeros((N, 4, 4))
    T[:, 0:3, 0:3] = R.from_quat(position_quat[:, 3:]).as_matrix()
    T[:, :3, 3] = position_quat[:, :3]
    T[:, 3, 3] = 1
    # except ValueError:
    #     print("Zero quat error!")
    # return T.squeeze()
    return T


def tf2isaacgympose(cam_pose):
    r = R.from_matrix(cam_pose[:3, :3])
    axis_angle = r.as_rotvec()
    norm = np.linalg.norm(axis_angle) 
    transform = gymapi.Transform()
    transform.p = gymapi.Vec3(
        cam_pose[0, 3], cam_pose[1, 3], cam_pose[2, 3]
    )
    transform.r = gymapi.Quat.from_axis_angle(
        gymapi.Vec3(axis_angle[0], axis_angle[1], axis_angle[2]), norm
    )
    return transform

def flip_camera_z(cam_pose):
    '''
    flip camera pose by 180 deg in the X-axis, so depth is negative Z
    '''
    quat_real_to_tac_neural = gymapi.Quat.from_euler_zyx(np.radians(180), 0.0, 0.0)
    flip_cam_pose = copy.deepcopy(cam_pose)
    flip_cam_pose.r = cam_pose.r * quat_real_to_tac_neural
    return flip_cam_pose

def isaacgym2tacneural(cam_pose):
    cam_pose = np.array(
        [
            cam_pose.p.x,
            cam_pose.p.y,
            cam_pose.p.z,
            cam_pose.r.x,
            cam_pose.r.y,
            cam_pose.r.z,
            cam_pose.r.w,
        ]
    )
    return xyzquat_to_tf_numpy(cam_pose)
    
def xyzw_to_wxyz(quat: torch.Tensor) -> torch.Tensor:
    """
    Convention change: [x, y, z, qx, qy, qz, qw] --> [x, y, z, qw, qx, qy, qz]
    """
    if quat.shape[1] == 7:
        return quat[:, [0, 1, 2, 6, 3, 4, 5]]
    else:
        return quat[:, [3, 0, 1, 2]]


def wxyz_to_xyzw(quat: torch.Tensor) -> torch.Tensor:
    """
    Convention change: [x, y, z, qw, qx, qy, qz] --> [x, y, z, qx, qy, qz, qw]
    """
    if quat.shape[1] == 7:
        return quat[:, [0, 1, 2, 4, 5, 6, 3]]
    else:
        return quat[:, [1, 2, 3, 0]]


def transform_pc(pointclouds: np.ndarray, poses: np.ndarray):
    """
    Transform pointclouds by poses
    """

    if type(pointclouds) is not list:
        temp = pointclouds
        pointclouds = [None] * 1
        pointclouds[0] = temp
        poses = np.expand_dims(poses, axis=2)

    if len(poses.shape) < 3:
        poses = xyzquat_to_tf_numpy(poses)

    transformed_pointclouds = pointclouds
    # TODO: vectorize
    for i, (pointcloud, pose) in enumerate(zip(pointclouds, poses)):
        pointcloud = pointcloud.T
        # 3D affine transform
        pointcloud = pose @ np.vstack([pointcloud, np.ones((1, pointcloud.shape[1]))])
        pointcloud = pointcloud / pointcloud[3, :]
        pointcloud = pointcloud[:3, :].T
        transformed_pointclouds[i] = pointcloud
    return (
        transformed_pointclouds[0] if len(pointclouds) == 1 else transformed_pointclouds
    )


def wrap_angles(angles: torch.Tensor) -> torch.Tensor:
    """
    angles : (N, 3) angles in degrees
    Wraps to [-np.pi, np.pi] or [-180, 180]
    """

    mask = angles > 180.0
    angles[mask] -= 2.0 * 180.0

    mask = angles < -180.0
    angles[mask] += 2.0 * 180.0
    return angles


def quat2euler(quat: np.ndarray) -> np.ndarray:
    """
    Convert quaternions to euler angles
    """
    quat = np.atleast_2d(quat)
    r = R.from_quat(quat)
    return r.as_euler("xyz", degrees=True)


def rot2euler(rot: torch.Tensor) -> torch.Tensor:
    """
    Convert rotation matrix to euler angles
    Adapted from so3_rotation_angle() in  https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/so3.html
    """
    rot_trace = rot[:, 0, 0] + rot[:, 1, 1] + rot[:, 2, 2]
    phi_cos = torch.acos((rot_trace - 1.0) * 0.5)
    return torch.rad2deg(phi_cos)

    # r = R.from_matrix(np.atleast_3d(rot.cpu().numpy()))
    # eul = r.as_euler('xyz', degrees = True)
    # return torch.tensor(eul, device= rot.device)


def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
    """
    https://pytorch3d.readthedocs.io/en/latest/modules/transforms.html#pytorch3d.transforms.euler_angles_to_matrix
    Convert rotations given as Euler angles in radians to rotation matrices.
    Args:
        euler_angles: Euler angles in radians as tensor of shape (..., 3).
        convention: Convention string of three uppercase letters from
            {"X", "Y", and "Z"}.
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
        raise ValueError("Invalid input euler angles.")
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    matrices = [
        _axis_angle_rotation(c, e)
        for c, e in zip(convention, torch.unbind(euler_angles, -1))
    ]
    # return functools.reduce(torch.matmul, matrices)
    return matrices[0] @ matrices[1] @ matrices[2]


def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
    """
    https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html
    Return the rotation matrices for one of the rotations about an axis
    of which Euler angles describe, for each value of the angle given.
    Args:
        axis: Axis label "X" or "Y or "Z".
        angle: any shape tensor of Euler angles in radians
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """

    cos = torch.cos(angle)
    sin = torch.sin(angle)
    one = torch.ones_like(angle)
    zero = torch.zeros_like(angle)

    if axis == "X":
        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
    elif axis == "Y":
        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
    elif axis == "Z":
        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
    else:
        raise ValueError("letter must be either X, Y or Z.")

    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))

def skew_matrix(v: np.ndarray) -> np.ndarray:
    """
    Get skew-symmetric matrix from vector
    """
    v = np.atleast_2d(v)
    # vector to its skew matrix
    mat = np.zeros((3, 3, v.shape[0]))
    mat[0, 1, :] = -1 * v[:, 2]
    mat[0, 2, :] = v[:, 1]

    mat[1, 0, :] = v[:, 2]
    mat[1, 2, :] = -1 * v[:, 0]

    mat[2, 0, :] = -1 * v[:, 1]
    mat[2, 1, :] = v[:, 0]
    return mat


def pose_from_vertex_normal(
    vertices: np.ndarray, normals: np.ndarray, shear_mag: float, delta: np.ndarray
) -> np.ndarray:
    """
    Generate SE(3) pose given
    vertices: (N, 3), normals: (N, 3), shear_mag: scalar, delta: (N, 1)
    """
    vertices = np.atleast_2d(vertices)
    normals = np.atleast_2d(normals)

    num_samples = vertices.shape[0]
    T = np.zeros((num_samples, 4, 4))  # transform from point coord to world coord
    T[:, 3, 3] = 1
    T[:, :3, 3] = vertices  # t

    # resolve ambiguous DoF
    """Find rotation of shear_vector so its orientation matches normal: np.dot(Rot, shear_vector) = normal
    https://math.stackexchange.com/questions/293116/rotating-one-3d-vector-to-another """

    cos_shear_mag = np.random.uniform(
        low=np.cos(shear_mag), high=1.0, size=(num_samples,)
    )  # Base of shear cone
    shear_phi = np.random.uniform(
        low=0.0, high=2 * np.pi, size=(num_samples,)
    )  # Circle of shear cone

    # Axis v = (shear_vector \cross normal)/(||shear_vector \cross normal||)
    shear_vector = np.array(
        [
            np.sqrt(1 - cos_shear_mag**2) * np.cos(shear_phi),
            np.sqrt(1 - cos_shear_mag**2) * np.sin(shear_phi),
            cos_shear_mag,
        ]
    ).T
    shear_vector_skew = skew_matrix(shear_vector)
    v = np.einsum("ijk,jk->ik", shear_vector_skew, normals.T).T
    v = v / np.linalg.norm(v, axis=1).reshape(-1, 1)

    # find corner cases
    check = np.einsum("ij,ij->i", normals, np.array([[0, 0, 1]]))
    zero_idx_up = check > 0.9  # pointing up
    zero_idx_down = check < -0.9  # pointing down

    v_skew, sampledNormals_skew = skew_matrix(v), skew_matrix(normals)

    # Angle theta = \arccos(z_axis \dot normal)
    # elementwise: theta = np.arccos(np.dot(shear_vector,normal)/(np.linalg.norm(shear_vector)*np.linalg.norm(normal)))
    theta = np.arccos(
        np.einsum("ij,ij->i", shear_vector, normals)
        / (np.linalg.norm(shear_vector, axis=1) * np.linalg.norm(normals, axis=1))
    )

    identity_3d = np.zeros(v_skew.shape)
    np.einsum("iij->ij", identity_3d)[:] = 1
    # elementwise: Rot = np.identity(3) + v_skew*np.sin(theta) + np.linalg.matrix_power(v_skew,2) * (1-np.cos(theta)) # rodrigues
    Rot = (
        identity_3d
        + v_skew * np.sin(theta)
        + np.einsum("ijn,jkn->ikn", v_skew, v_skew) * (1 - np.cos(theta))
    )  # rodrigues

    if np.any(zero_idx_up):
        Rot[:3, :3, zero_idx_up] = np.dstack([np.identity(3)] * np.sum(zero_idx_up))
    if np.any(zero_idx_down):
        Rot[:3, :3, zero_idx_down] = np.dstack(
            [np.array([[1, 0, 0], [0, -1, -0], [0, 0, -1]])] * np.sum(zero_idx_down)
        )

    # Rotation about Z axis is still ambiguous, generating random rotation b/w [0, 2pi] about normal axis
    # elementwise: RotDelta = np.identity(3) + normal_skew*np.sin(delta[i]) + np.linalg.matrix_power(normal_skew,2) * (1-np.cos(delta[i])) # rodrigues
    RotDelta = (
        identity_3d
        + sampledNormals_skew * np.sin(delta)
        + np.einsum("ijn,jkn->ikn", sampledNormals_skew, sampledNormals_skew)
        * (1 - np.cos(delta))
    )  # rodrigues

    # elementwise:  RotDelta @ Rot
    tfs = np.einsum("ijn,jkn->ikn", RotDelta, Rot)
    T[:, :3, :3] = np.rollaxis(tfs, 2)
    return T


def cam2gel(cam_pose, cam_dist):
    """
    Convert cam_pose to gel_pose
    """
    cam_tf = torch.eye(4, device=cam_pose.device)
    cam_tf[2, 3] = -cam_dist
    return cam_pose @ cam_tf[None, :]
NCF_policies/isaacgymenvs/tasks/ncf_tacto/pose.py
carolinahiguera/NCF-Policy
import argparse
import logging
import os
from datetime import datetime
import threading
from typing import Optional

# noinspection PyUnresolvedReferences
import isaacgym

import hydra
from termcolor import cprint

from isaacgymenvs.utils.rlgames_utils import multi_gpu_get_rank
from omegaconf import DictConfig, OmegaConf
from hydra.utils import to_absolute_path
from isaacgymenvs.tasks import isaacgym_task_map
from omegaconf import DictConfig, OmegaConf
import gym

from isaacgymenvs.utils.reformat import omegaconf_to_dict, print_dict
from isaacgymenvs.utils.utils import set_np_formatting, set_seed

from algo2.policy.ppo.ppo import PPO as PPO2
import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)


# @hydra.main(config_name="config", config_path="./cfg")
def run(cfg: DictConfig, config_path: Optional[str] = None):
    # hydra.initialize(config_path="./cfg", job_name="train")

    if config_path is not None:
        assert cfg is None
        cfg = OmegaConf.load(config_path)

    if cfg.checkpoint:
        if "*" in cfg.checkpoint:
            from glob import glob

            _ckpt = glob(cfg.checkpoint)
            assert len(_ckpt) == 1
            cfg.checkpoint = _ckpt[0]
        cfg.checkpoint = to_absolute_path(cfg.checkpoint)

    # set numpy formatting for printing only
    set_np_formatting()

    # global rank of the GPU
    if cfg.train.ppo.multi_gpu:
        rank = int(os.getenv("LOCAL_RANK", "0"))
        # torchrun --standalone --nnodes=1 --nproc_per_node=2 train.py
        cfg.sim_device = f"cuda:{rank}"
        cfg.rl_device = f"cuda:{rank}"
        cfg.graphics_device_id = int(rank)
        # sets seed. if seed is -1 will pick a random one
        cfg.seed = set_seed(cfg.seed + rank)
    else:
        rank = -1
        cfg.seed = set_seed(cfg.seed)

    cprint("Start Building the Environment", "green", attrs=["bold"])

    envs = isaacgym_task_map[cfg.task_name](
        cfg=omegaconf_to_dict(cfg.task),
        sim_device=cfg.sim_device,
        rl_device=cfg.rl_device,
        # graphics_device_id=-1,  # if cfg.headless else cfg.graphics_device_id,
        graphics_device_id=0,
        # graphics_device_id=cfg.graphics_device_id,
        headless=cfg.headless,
        virtual_screen_capture=False,  # cfg.capture_video,
        force_render=cfg.force_render if not cfg.headless else False,
    )

    output_dif = os.path.join(cfg.outputs_root_dir, cfg.train.ppo.output_name)
    os.makedirs(output_dif, exist_ok=True)
    agent = eval(cfg.train.algo)(envs, output_dif, full_config=cfg)

    if cfg.test:
        assert cfg.train.load_path
        agent.restore_test(cfg.train.load_path)
        agent.test()
    else:
        if rank <= 0:
            date = str(datetime.now().strftime("%m%d%H"))
            with open(os.path.join(output_dif, f"config_{date}.yaml"), "w") as f:
                f.write(OmegaConf.to_yaml(cfg))

        # check whether execute train by mistake:
        best_ckpt_path = os.path.join(
            "outputs",
            cfg.train.ppo.output_name,
            "stage1_nn" if cfg.train.algo == "PPO" else "stage2_nn",
            "best.pth",
        )
        if os.path.exists(best_ckpt_path):
            user_input = input(
                f"are you intentionally going to overwrite files in {cfg.train.ppo.output_name}, type yes to continue \n"
            )
            if user_input != "yes":
                exit()

        agent.restore_train(cfg.train.load_path)
        agent.train()

    cprint("Finished", "green", attrs=["bold"])


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "config_path", type=argparse.FileType("r"), help="Path to hydra config."
    )
    args = parser.parse_args()
    run(None, args.config_path)

NCF_policies/isaacgymenvs/train.py
carolinahiguera/NCF-Policy
import pathlib
import sys
import tempfile
import traceback

import hydra
import submitit
from omegaconf import DictConfig, OmegaConf
from hydra.core.hydra_config import HydraConfig

import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "1"


def launch_submitit(config):
    launcher_cfg = HydraConfig.get().launcher
    assert launcher_cfg.tasks_per_node == 1 and launcher_cfg.nodes == 1
    with tempfile.NamedTemporaryFile() as fp:
        OmegaConf.save(config=config, f=fp.name)
        if config.train.ppo.multi_gpu:
            exec_cmd = [
                "torchrun",
                "--standalone",
                "--nnodes=1",
                f"--nproc_per_node={launcher_cfg.gpus_per_node}",
            ]
        else:
            exec_cmd = ["python3.8"]
        function = submitit.helpers.CommandFunction(
            exec_cmd + ["NCF_policies/isaacgymenvs/train.py", f"{fp.name}"]
        )
        executor = submitit.AutoExecutor(
            folder=pathlib.Path(launcher_cfg.submitit_folder).parent / "train_logs",
            cluster="local",
        )
        executor.update_parameters(
            timeout_min=launcher_cfg.timeout_min,
            cpus_per_task=launcher_cfg.cpus_per_task,
            gpus_per_node=launcher_cfg.gpus_per_node,
            mem_gb=launcher_cfg.mem_gb,
        )
        job = executor.submit(function)
        print(job.result())


@hydra.main(config_name="config", config_path="./cfg")
def main(config: DictConfig):
    try:
        launcher_cfg = HydraConfig.get().launcher
        if "submitit" in launcher_cfg._target_:
            launch_submitit(config)
        else:
            # lazy import so that submitit launch is faster
            # from isaacgymenvs.train import run
            from train import run

            run(config)

    except Exception:
        traceback.print_exc(file=sys.stderr)
        raise


if __name__ == "__main__":
    main()

NCF_policies/isaacgymenvs/train_cluster.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import numpy as np
from bisect import bisect
from isaacgym import gymapi


def get_property_setter_map(gym):
    property_to_setters = {
        "dof_properties": gym.set_actor_dof_properties,
        "tendon_properties": gym.set_actor_tendon_properties,
        "rigid_body_properties": gym.set_actor_rigid_body_properties,
        "rigid_shape_properties": gym.set_actor_rigid_shape_properties,
        "sim_params": gym.set_sim_params,
    }

    return property_to_setters


def get_property_getter_map(gym):
    property_to_getters = {
        "dof_properties": gym.get_actor_dof_properties,
        "tendon_properties": gym.get_actor_tendon_properties,
        "rigid_body_properties": gym.get_actor_rigid_body_properties,
        "rigid_shape_properties": gym.get_actor_rigid_shape_properties,
        "sim_params": gym.get_sim_params,
    }

    return property_to_getters


def get_default_setter_args(gym):
    property_to_setter_args = {
        "dof_properties": [],
        "tendon_properties": [],
        "rigid_body_properties": [True],
        "rigid_shape_properties": [],
        "sim_params": [],
    }

    return property_to_setter_args


def generate_random_samples(attr_randomization_params, shape, curr_gym_step_count,
                            extern_sample=None):

    rand_range = attr_randomization_params['range']
    distribution = attr_randomization_params['distribution']

    sched_type = attr_randomization_params['schedule'] if 'schedule' in attr_randomization_params else None
    sched_step = attr_randomization_params['schedule_steps'] if 'schedule' in attr_randomization_params else None

    operation = attr_randomization_params['operation']

    if sched_type == 'linear':
        sched_scaling = 1 / sched_step * min(curr_gym_step_count, sched_step)
    elif sched_type == 'constant':
        sched_scaling = 0 if curr_gym_step_count < sched_step else 1
    else:
        sched_scaling = 1

    if extern_sample is not None:

        sample = extern_sample

        if operation == 'additive':
            sample *= sched_scaling
        elif operation == 'scaling':
            sample = sample * sched_scaling + 1 * (1 - sched_scaling)

    elif distribution == "gaussian":

        mu, var = rand_range

        if operation == 'additive':
            mu *= sched_scaling
            var *= sched_scaling
        elif operation == 'scaling':
            var = var * sched_scaling  # scale up var over time
            mu = mu * sched_scaling + 1 * (1 - sched_scaling)  # linearly interpolate
        sample = np.random.normal(mu, var, shape)

    elif distribution == "loguniform":

        lo, hi = rand_range
        if operation == 'additive':
            lo *= sched_scaling
            hi *= sched_scaling
        elif operation == 'scaling':
            lo = lo * sched_scaling + 1 * (1 - sched_scaling)
            hi = hi * sched_scaling + 1 * (1 - sched_scaling)
        sample = np.exp(np.random.uniform(np.log(lo), np.log(hi), shape))

    elif distribution == "uniform":

        lo, hi = rand_range
        if operation == 'additive':
            lo *= sched_scaling
            hi *= sched_scaling
        elif operation == 'scaling':
            lo = lo * sched_scaling + 1 * (1 - sched_scaling)
            hi = hi * sched_scaling + 1 * (1 - sched_scaling)
        sample = np.random.uniform(lo, hi, shape)

    return sample


def get_bucketed_val(new_prop_val, attr_randomization_params):
    if attr_randomization_params['distribution'] == 'uniform':
        # range of buckets defined by uniform distribution
        lo, hi = attr_randomization_params['range'][0], attr_randomization_params['range'][1]
    else:
        # for gaussian, set range of buckets to be 2 stddev away from mean
        lo = attr_randomization_params['range'][0] - 2 * np.sqrt(attr_randomization_params['range'][1])
        hi = attr_randomization_params['range'][0] + 2 * np.sqrt(attr_randomization_params['range'][1])
    num_buckets = attr_randomization_params['num_buckets']
    buckets = [(hi - lo) * i / num_buckets + lo for i in range(num_buckets)]
    return buckets[bisect(buckets, new_prop_val) - 1]


def apply_random_samples(prop, og_prop, attr, attr_randomization_params,
                         curr_gym_step_count, extern_sample=None, bucketing_randomization_params=None):
    
    """
    @params:
        prop: property we want to randomise
        og_prop: the original property and its value 
        attr: which particular attribute we want to randomise e.g. damping, stiffness
        attr_randomization_params: the attribute randomisation meta-data e.g. distr, range, schedule
        curr_gym_step_count: gym steps so far 

    """

    if isinstance(prop, gymapi.SimParams):
        
        if attr == 'gravity':
            
            sample = generate_random_samples(attr_randomization_params, 3, curr_gym_step_count)
            if attr_randomization_params['operation'] == 'scaling':
                prop.gravity.x = og_prop['gravity'].x * sample[0]
                prop.gravity.y = og_prop['gravity'].y * sample[1]
                prop.gravity.z = og_prop['gravity'].z * sample[2]
            elif attr_randomization_params['operation'] == 'additive':
                prop.gravity.x = og_prop['gravity'].x + sample[0]
                prop.gravity.y = og_prop['gravity'].y + sample[1]
                prop.gravity.z = og_prop['gravity'].z + sample[2]


        if attr == 'rest_offset':

           sample = generate_random_samples(attr_randomization_params, 1, curr_gym_step_count)
           prop.physx.rest_offset = sample
                

    elif isinstance(prop, np.ndarray):
        sample = generate_random_samples(attr_randomization_params, prop[attr].shape,
                                         curr_gym_step_count, extern_sample)

        if attr_randomization_params['operation'] == 'scaling':
            new_prop_val = og_prop[attr] * sample
        elif attr_randomization_params['operation'] == 'additive':
            new_prop_val = og_prop[attr] + sample

        if 'num_buckets' in attr_randomization_params and attr_randomization_params['num_buckets'] > 0:
            new_prop_val = get_bucketed_val(new_prop_val, attr_randomization_params)
        prop[attr] = new_prop_val
    else:
        sample = generate_random_samples(attr_randomization_params, 1,
                                         curr_gym_step_count, extern_sample)
        cur_attr_val = og_prop[attr]
        if attr_randomization_params['operation'] == 'scaling':
            new_prop_val = cur_attr_val * sample
        elif attr_randomization_params['operation'] == 'additive':
            new_prop_val = cur_attr_val + sample

        if 'num_buckets' in attr_randomization_params and attr_randomization_params['num_buckets'] > 0:
            if bucketing_randomization_params is None:
                new_prop_val = get_bucketed_val(new_prop_val, attr_randomization_params)
            else:
                new_prop_val = get_bucketed_val(new_prop_val, bucketing_randomization_params)
        setattr(prop, attr, new_prop_val)

def check_buckets(gym, envs, dr_params):
    total_num_buckets = 0
    for actor, actor_properties in dr_params["actor_params"].items():
        cur_num_buckets = 0

        if 'rigid_shape_properties' in actor_properties.keys():
            prop_attrs = actor_properties['rigid_shape_properties']
            if 'restitution' in prop_attrs and 'num_buckets' in prop_attrs['restitution']:
                cur_num_buckets = prop_attrs['restitution']['num_buckets']
            if 'friction' in prop_attrs and 'num_buckets' in prop_attrs['friction']:
                if cur_num_buckets > 0:
                    cur_num_buckets *= prop_attrs['friction']['num_buckets']
                else:
                    cur_num_buckets = prop_attrs['friction']['num_buckets']
            total_num_buckets += cur_num_buckets

    assert total_num_buckets <= 64000, 'Explicit material bucketing has been specified, but the provided total bucket count exceeds 64K: {} specified buckets'.format(
        total_num_buckets)

    shape_ct = 0

    # Separate loop because we should not assume that each actor is present in each env
    for env in envs:
        for i in range(gym.get_actor_count(env)):
            actor_handle = gym.get_actor_handle(env, i)
            actor_name = gym.get_actor_name(env, actor_handle)
            if actor_name in dr_params["actor_params"] and 'rigid_shape_properties' in dr_params["actor_params"][actor_name]:
                shape_ct += gym.get_actor_rigid_shape_count(env, actor_handle)

    assert shape_ct <= 64000 or total_num_buckets > 0, 'Explicit material bucketing is not used but the total number of shapes exceeds material limit. Please specify bucketing to limit material count.'
NCF_policies/isaacgymenvs/utils/dr_utils.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
from collections import deque
from typing import Callable, Dict, Tuple, Any

import os
import gym
import numpy as np
import torch
from rl_games.common import env_configurations, vecenv
from rl_games.common.algo_observer import AlgoObserver

from isaacgymenvs.tasks import isaacgym_task_map
from isaacgymenvs.utils.utils import set_seed, flatten_dict


def multi_gpu_get_rank(multi_gpu):
    if multi_gpu:
        rank = int(os.getenv("LOCAL_RANK", "0"))
        print("GPU rank: ", rank)
        return rank

    return 0


def get_rlgames_env_creator(
        # used to create the vec task
        seed: int,
        task_config: dict,
        task_name: str,
        sim_device: str,
        rl_device: str,
        graphics_device_id: int,
        headless: bool,
        # used to handle multi-gpu case
        multi_gpu: bool = False,
        post_create_hook: Callable = None,
        virtual_screen_capture: bool = False,
        force_render: bool = False,
):
    """Parses the configuration parameters for the environment task and creates a VecTask

    Args:
        task_config: environment configuration.
        task_name: Name of the task, used to evaluate based on the imported name (eg 'Trifinger')
        sim_device: The type of env device, eg 'cuda:0'
        rl_device: Device that RL will be done on, eg 'cuda:0'
        graphics_device_id: Graphics device ID.
        headless: Whether to run in headless mode.
        multi_gpu: Whether to use multi gpu
        post_create_hook: Hooks to be called after environment creation.
            [Needed to setup WandB only for one of the RL Games instances when doing multiple GPUs]
        virtual_screen_capture: Set to True to allow the users get captured screen in RGB array via `env.render(mode='rgb_array')`. 
        force_render: Set to True to always force rendering in the steps (if the `control_freq_inv` is greater than 1 we suggest stting this arg to True)
    Returns:
        A VecTaskPython object.
    """
    def create_rlgpu_env():
        """
        Creates the task from configurations and wraps it using RL-games wrappers if required.
        """
        if multi_gpu:

            local_rank = int(os.getenv("LOCAL_RANK", "0"))
            global_rank = int(os.getenv("RANK", "0"))

            # local rank of the GPU in a node
            local_rank = int(os.getenv("LOCAL_RANK", "0"))
            # global rank of the GPU
            global_rank = int(os.getenv("RANK", "0"))
            # total number of GPUs across all nodes
            world_size = int(os.getenv("WORLD_SIZE", "1"))

            print(f"global_rank = {global_rank} local_rank = {local_rank} world_size = {world_size}")

            _sim_device = f'cuda:{local_rank}'
            _rl_device = f'cuda:{local_rank}'

            task_config['rank'] = local_rank
            task_config['rl_device'] = _rl_device
        else:
            _sim_device = sim_device
            _rl_device = rl_device

        # create native task and pass custom config
        env = isaacgym_task_map[task_name](
            cfg=task_config,
            rl_device=_rl_device,
            sim_device=_sim_device,
            graphics_device_id=graphics_device_id,
            headless=headless,
            virtual_screen_capture=virtual_screen_capture,
            force_render=force_render,
        )

        if post_create_hook is not None:
            post_create_hook()

        return env
    return create_rlgpu_env


class RLGPUAlgoObserver(AlgoObserver):
    """Allows us to log stats from the env along with the algorithm running stats. """

    def __init__(self):
        super().__init__()
        self.algo = None
        self.writer = None

        self.ep_infos = []
        self.direct_info = {}

        self.episode_cumulative = dict()
        self.episode_cumulative_avg = dict()
        self.new_finished_episodes = False

    def after_init(self, algo):
        self.algo = algo
        self.writer = self.algo.writer

    def process_infos(self, infos, done_indices):
        assert isinstance(infos, dict), 'RLGPUAlgoObserver expects dict info'
        if not isinstance(infos, dict):
            return

        if 'episode' in infos:
            self.ep_infos.append(infos['episode'])

        if 'episode_cumulative' in infos:
            for key, value in infos['episode_cumulative'].items():
                if key not in self.episode_cumulative:
                    self.episode_cumulative[key] = torch.zeros_like(value)
                self.episode_cumulative[key] += value

            for done_idx in done_indices:
                self.new_finished_episodes = True
                done_idx = done_idx.item()

                for key, value in infos['episode_cumulative'].items():
                    if key not in self.episode_cumulative_avg:
                        self.episode_cumulative_avg[key] = deque([], maxlen=self.algo.games_to_track)

                    self.episode_cumulative_avg[key].append(self.episode_cumulative[key][done_idx].item())
                    self.episode_cumulative[key][done_idx] = 0

        # turn nested infos into summary keys (i.e. infos['scalars']['lr'] -> infos['scalars/lr']
        if len(infos) > 0 and isinstance(infos, dict):  # allow direct logging from env
            infos_flat = flatten_dict(infos, prefix='', separator='/')
            self.direct_info = {}
            for k, v in infos_flat.items():
                # only log scalars
                if isinstance(v, float) or isinstance(v, int) or (isinstance(v, torch.Tensor) and len(v.shape) == 0):
                    self.direct_info[k] = v

    def after_print_stats(self, frame, epoch_num, total_time):
        if self.ep_infos:
            for key in self.ep_infos[0]:
                infotensor = torch.tensor([], device=self.algo.device)
                for ep_info in self.ep_infos:
                    # handle scalar and zero dimensional tensor infos
                    if not isinstance(ep_info[key], torch.Tensor):
                        ep_info[key] = torch.Tensor([ep_info[key]])
                    if len(ep_info[key].shape) == 0:
                        ep_info[key] = ep_info[key].unsqueeze(0)
                    infotensor = torch.cat((infotensor, ep_info[key].to(self.algo.device)))
                value = torch.mean(infotensor)
                self.writer.add_scalar('Episode/' + key, value, epoch_num)
            self.ep_infos.clear()
        
        # log these if and only if we have new finished episodes
        if self.new_finished_episodes:
            for key in self.episode_cumulative_avg:
                self.writer.add_scalar(f'episode_cumulative/{key}', np.mean(self.episode_cumulative_avg[key]), frame)
                self.writer.add_scalar(f'episode_cumulative_min/{key}_min', np.min(self.episode_cumulative_avg[key]), frame)
                self.writer.add_scalar(f'episode_cumulative_max/{key}_max', np.max(self.episode_cumulative_avg[key]), frame)
            self.new_finished_episodes = False

        for k, v in self.direct_info.items():
            self.writer.add_scalar(f'{k}/frame', v, frame)
            self.writer.add_scalar(f'{k}/iter', v, epoch_num)
            self.writer.add_scalar(f'{k}/time', v, total_time)


class MultiObserver(AlgoObserver):
    """Meta-observer that allows the user to add several observers."""

    def __init__(self, observers_):
        super().__init__()
        self.observers = observers_

    def _call_multi(self, method, *args_, **kwargs_):
        for o in self.observers:
            getattr(o, method)(*args_, **kwargs_)

    def before_init(self, base_name, config, experiment_name):
        self._call_multi('before_init', base_name, config, experiment_name)

    def after_init(self, algo):
        self._call_multi('after_init', algo)

    def process_infos(self, infos, done_indices):
        self._call_multi('process_infos', infos, done_indices)

    def after_steps(self):
        self._call_multi('after_steps')

    def after_clear_stats(self):
        self._call_multi('after_clear_stats')

    def after_print_stats(self, frame, epoch_num, total_time):
        self._call_multi('after_print_stats', frame, epoch_num, total_time)


class RLGPUEnv(vecenv.IVecEnv):
    def __init__(self, config_name, num_actors, **kwargs):
        self.env = env_configurations.configurations[config_name]['env_creator'](**kwargs)

    def step(self, actions):
        return self.env.step(actions)

    def reset(self):
        return self.env.reset()
    
    def reset_done(self):
        return self.env.reset_done()

    def get_number_of_agents(self):
        return self.env.get_number_of_agents()

    def get_env_info(self):
        info = {}
        info['action_space'] = self.env.action_space
        info['observation_space'] = self.env.observation_space

        if hasattr(self.env, "amp_observation_space"):
            info['amp_observation_space'] = self.env.amp_observation_space

        if self.env.num_states > 0:
            info['state_space'] = self.env.state_space
            print(info['action_space'], info['observation_space'], info['state_space'])
        else:
            print(info['action_space'], info['observation_space'])

        return info

    def set_train_info(self, env_frames, *args_, **kwargs_):
        """
        Send the information in the direction algo->environment.
        Most common use case: tell the environment how far along we are in the training process. This is useful
        for implementing curriculums and things such as that.
        """
        if hasattr(self.env, 'set_train_info'):
            self.env.set_train_info(env_frames, *args_, **kwargs_)

    def get_env_state(self):
        """
        Return serializable environment state to be saved to checkpoint.
        Can be used for stateful training sessions, i.e. with adaptive curriculums.
        """
        if hasattr(self.env, 'get_env_state'):
            return self.env.get_env_state()
        else:
            return None

    def set_env_state(self, env_state):
        if hasattr(self.env, 'set_env_state'):
            self.env.set_env_state(env_state)


class ComplexObsRLGPUEnv(vecenv.IVecEnv):
    
    def __init__(
        self,
        config_name,
        num_actors,
        obs_spec: Dict[str, Dict],
        **kwargs,
    ):
        """RLGPU wrapper for Isaac Gym tasks.

        Args:
            config_name: Name of rl games env_configurations configuration to use.
            obs_spec: Dictinoary listing out specification for observations to use.
                eg.
                {
                 'obs': {'names': ['obs_1', 'obs_2'], 'concat': True, space_name: 'observation_space'},},
                 'states': {'names': ['state_1', 'state_2'], 'concat': False, space_name: 'state_space'},}
                }
                Within each, if 'concat' is set, concatenates all the given observaitons into a single tensor of dim (num_envs, sum(num_obs)).
                    Assumes that each indivdual observation is single dimensional (ie (num_envs, k), so image observation isn't supported).
                    Currently applies to student and teacher both.
                "space_name" is given into the env info which RL Games reads to find the space shape
        """
        self.env = env_configurations.configurations[config_name]['env_creator'](**kwargs)

        self.obs_spec = obs_spec

    def _generate_obs(
        self, env_obs: Dict[str, torch.Tensor]
    ) -> Dict[str, Dict[str, torch.Tensor]]:
        """Generate the RL Games observations given the observations from the environment.

        Args:
            env_obs: environment observations
        Returns:
            Dict which contains keys with values corresponding to observations.
        """
        # rl games expects a dictionary with 'obs' and 'states'
        # corresponding to the policy observations and possible asymmetric
        # observations respectively

        rlgames_obs = {k: self.gen_obs_dict(env_obs, v['names'], v['concat']) for k, v in self.obs_spec.items()}

        return rlgames_obs

    def step(
        self, action: torch.Tensor
    ) -> Tuple[
        Dict[str, Dict[str, torch.Tensor]], torch.Tensor, torch.Tensor, Dict[str, Any]
    ]:
        """Step the Isaac Gym task.

        Args:
            action: Enivronment action.
        Returns:
            observations, rewards, dones, infos
            Returned obeservations are a dict which contains key 'obs' corresponding to a dictionary of observations,
            and possible 'states' key corresponding to dictionary of privileged observations.
        """
        env_obs, rewards, dones, infos = self.env.step(action)
        rlgames_obs = self._generate_obs(env_obs)
        return rlgames_obs, rewards, dones, infos

    def reset(self) -> Dict[str, Dict[str, torch.Tensor]]:
        env_obs = self.env.reset()
        return self._generate_obs(env_obs)

    def get_number_of_agents(self) -> int:
        return self.env.get_number_of_agents()

    def get_env_info(self) -> Dict[str, gym.spaces.Space]:
        """Gets information on the environment's observation, action, and privileged observation (states) spaces."""
        info = {}
        info["action_space"] = self.env.action_space

        for k, v in self.obs_spec.items():
            info[v['space_name']] = self.gen_obs_space(v['names'], v['concat'])

        return info
    
    def gen_obs_dict(self, obs_dict, obs_names, concat):
        """Generate the RL Games observations given the observations from the environment."""
        if concat:
            return torch.cat([obs_dict[name] for name in obs_names], dim=1)
        else:
            return {k: obs_dict[k] for k in obs_names}
            

    def gen_obs_space(self, obs_names, concat):
        """Generate the RL Games observation space given the observations from the environment."""
        if concat:
            return gym.spaces.Box(
                low=-np.Inf,
                high=np.Inf,
                shape=(sum([self.env.observation_space[s].shape[0] for s in obs_names]),),
                dtype=np.float32,
            )
        else:        
            return gym.spaces.Dict(
                    {k: self.env.observation_space[k] for k in obs_names}
                )

    def set_train_info(self, env_frames, *args_, **kwargs_):
        """
        Send the information in the direction algo->environment.
        Most common use case: tell the environment how far along we are in the training process. This is useful
        for implementing curriculums and things such as that.
        """
        if hasattr(self.env, 'set_train_info'):
            self.env.set_train_info(env_frames, *args_, **kwargs_)

    def get_env_state(self):
        """
        Return serializable environment state to be saved to checkpoint.
        Can be used for stateful training sessions, i.e. with adaptive curriculums.
        """
        if hasattr(self.env, 'get_env_state'):
            return self.env.get_env_state()
        else:
            return None

    def set_env_state(self, env_state):
        if hasattr(self.env, 'set_env_state'):
            self.env.set_env_state(env_state)                

NCF_policies/isaacgymenvs/utils/rlgames_utils.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import isaacgym

import torch
import torch.nn.functional as F
import numpy as np


def to_torch(x, dtype=torch.float, device='cuda:0', requires_grad=False):
    return torch.tensor(x, dtype=dtype, device=device, requires_grad=requires_grad)


@torch.jit.script
def quat_mul(a, b):
    assert a.shape == b.shape
    shape = a.shape
    a = a.reshape(-1, 4)
    b = b.reshape(-1, 4)

    x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3]
    x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
    ww = (z1 + x1) * (x2 + y2)
    yy = (w1 - y1) * (w2 + z2)
    zz = (w1 + y1) * (w2 - z2)
    xx = ww + yy + zz
    qq = 0.5 * (xx + (z1 - x1) * (x2 - y2))
    w = qq - ww + (z1 - y1) * (y2 - z2)
    x = qq - xx + (x1 + w1) * (x2 + w2)
    y = qq - yy + (w1 - x1) * (y2 + z2)
    z = qq - zz + (z1 + y1) * (w2 - x2)

    quat = torch.stack([x, y, z, w], dim=-1).view(shape)

    return quat


@torch.jit.script
def normalize(x, eps: float = 1e-9):
    return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)


@torch.jit.script
def quat_apply(a, b):
    shape = b.shape
    a = a.reshape(-1, 4)
    b = b.reshape(-1, 3)
    xyz = a[:, :3]
    t = xyz.cross(b, dim=-1) * 2
    return (b + a[:, 3:] * t + xyz.cross(t, dim=-1)).view(shape)


@torch.jit.script
def quat_rotate(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    a = v * (2.0 * q_w ** 2 - 1.0).unsqueeze(-1)
    b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
    c = q_vec * \
        torch.bmm(q_vec.view(shape[0], 1, 3), v.view(
            shape[0], 3, 1)).squeeze(-1) * 2.0
    return a + b + c


@torch.jit.script
def quat_rotate_inverse(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    a = v * (2.0 * q_w ** 2 - 1.0).unsqueeze(-1)
    b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
    c = q_vec * \
        torch.bmm(q_vec.view(shape[0], 1, 3), v.view(
            shape[0], 3, 1)).squeeze(-1) * 2.0
    return a - b + c


@torch.jit.script
def quat_conjugate(a):
    shape = a.shape
    a = a.reshape(-1, 4)
    return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape)


@torch.jit.script
def quat_unit(a):
    return normalize(a)


@torch.jit.script
def quat_from_angle_axis(angle, axis):
    theta = (angle / 2).unsqueeze(-1)
    xyz = normalize(axis) * theta.sin()
    w = theta.cos()
    return quat_unit(torch.cat([xyz, w], dim=-1))


@torch.jit.script
def normalize_angle(x):
    return torch.atan2(torch.sin(x), torch.cos(x))


@torch.jit.script
def tf_inverse(q, t):
    q_inv = quat_conjugate(q)
    return q_inv, -quat_apply(q_inv, t)


@torch.jit.script
def tf_apply(q, t, v):
    return quat_apply(q, v) + t


@torch.jit.script
def tf_vector(q, v):
    return quat_apply(q, v)


@torch.jit.script
def tf_combine(q1, t1, q2, t2):
    return quat_mul(q1, q2), quat_apply(q1, t2) + t1


@torch.jit.script
def get_basis_vector(q, v):
    return quat_rotate(q, v)


def get_axis_params(value, axis_idx, x_value=0., dtype=float, n_dims=3):
    """construct arguments to `Vec` according to axis index.
    """
    zs = np.zeros((n_dims,))
    assert axis_idx < n_dims, "the axis dim should be within the vector dimensions"
    zs[axis_idx] = 1.
    params = np.where(zs == 1., value, zs)
    params[0] = x_value
    return list(params.astype(dtype))


@torch.jit.script
def copysign(a, b):
    # type: (float, Tensor) -> Tensor
    a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
    return torch.abs(a) * torch.sign(b)


@torch.jit.script
def get_euler_xyz(q):
    qx, qy, qz, qw = 0, 1, 2, 3
    # roll (x-axis rotation)
    sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
    cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
        q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
    roll = torch.atan2(sinr_cosp, cosr_cosp)

    # pitch (y-axis rotation)
    sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
    pitch = torch.where(torch.abs(sinp) >= 1, copysign(
        np.pi / 2.0, sinp), torch.asin(sinp))

    # yaw (z-axis rotation)
    siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
    cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
        q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
    yaw = torch.atan2(siny_cosp, cosy_cosp)

    return roll % (2*np.pi), pitch % (2*np.pi), yaw % (2*np.pi)


@torch.jit.script
def quat_from_euler_xyz(roll, pitch, yaw):
    cy = torch.cos(yaw * 0.5)
    sy = torch.sin(yaw * 0.5)
    cr = torch.cos(roll * 0.5)
    sr = torch.sin(roll * 0.5)
    cp = torch.cos(pitch * 0.5)
    sp = torch.sin(pitch * 0.5)

    qw = cy * cr * cp + sy * sr * sp
    qx = cy * sr * cp - sy * cr * sp
    qy = cy * cr * sp + sy * sr * cp
    qz = sy * cr * cp - cy * sr * sp

    return torch.stack([qx, qy, qz, qw], dim=-1)


@torch.jit.script
def torch_rand_float(lower, upper, shape, device):
    # type: (float, float, Tuple[int, int], str) -> Tensor
    return (upper - lower) * torch.rand(*shape, device=device) + lower


@torch.jit.script
def torch_random_dir_2(shape, device):
    # type: (Tuple[int, int], str) -> Tensor
    angle = torch_rand_float(-np.pi, np.pi, shape, device).squeeze(-1)
    return torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)


@torch.jit.script
def tensor_clamp(t, min_t, max_t):
    return torch.max(torch.min(t, max_t), min_t)


@torch.jit.script
def scale(x, lower, upper):
    return (0.5 * (x + 1.0) * (upper - lower) + lower)


@torch.jit.script
def unscale(x, lower, upper):
    return (2.0 * x - upper - lower) / (upper - lower)


def unscale_np(x, lower, upper):
    return (2.0 * x - upper - lower) / (upper - lower)


@torch.jit.script
def compute_heading_and_up(
    torso_rotation, inv_start_rot, to_target, vec0, vec1, up_idx
):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]
    num_envs = torso_rotation.shape[0]
    target_dirs = normalize(to_target)

    torso_quat = quat_mul(torso_rotation, inv_start_rot)
    up_vec = get_basis_vector(torso_quat, vec1).view(num_envs, 3)
    heading_vec = get_basis_vector(torso_quat, vec0).view(num_envs, 3)
    up_proj = up_vec[:, up_idx]
    heading_proj = torch.bmm(heading_vec.view(
        num_envs, 1, 3), target_dirs.view(num_envs, 3, 1)).view(num_envs)

    return torso_quat, up_proj, heading_proj, up_vec, heading_vec


@torch.jit.script
def compute_rot(torso_quat, velocity, ang_velocity, targets, torso_positions):
    vel_loc = quat_rotate_inverse(torso_quat, velocity)
    angvel_loc = quat_rotate_inverse(torso_quat, ang_velocity)

    roll, pitch, yaw = get_euler_xyz(torso_quat)

    walk_target_angle = torch.atan2(targets[:, 2] - torso_positions[:, 2],
                                    targets[:, 0] - torso_positions[:, 0])
    angle_to_target = walk_target_angle - yaw

    return vel_loc, angvel_loc, roll, pitch, yaw, angle_to_target


@torch.jit.script
def quat_axis(q, axis=0):
    # type: (Tensor, int) -> Tensor
    basis_vec = torch.zeros(q.shape[0], 3, device=q.device)
    basis_vec[:, axis] = 1
    return quat_rotate(q, basis_vec)


"""
Normalization and Denormalization of Tensors
"""

@torch.jit.script
def scale_transform(x: torch.Tensor, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor:
    """
    Normalizes a given input tensor to a range of [-1, 1].

    @note It uses pytorch broadcasting functionality to deal with batched input.

    Args:
        x: Input tensor of shape (N, dims).
        lower: The minimum value of the tensor. Shape (dims,)
        upper: The maximum value of the tensor. Shape (dims,)

    Returns:
        Normalized transform of the tensor. Shape (N, dims)
    """
    # default value of center
    offset = (lower + upper) * 0.5
    # return normalized tensor
    return 2 * (x - offset) / (upper - lower)


@torch.jit.script
def unscale_transform(x: torch.Tensor, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor:
    """
    Denormalizes a given input tensor from range of [-1, 1] to (lower, upper).

    @note It uses pytorch broadcasting functionality to deal with batched input.

    Args:
        x: Input tensor of shape (N, dims).
        lower: The minimum value of the tensor. Shape (dims,)
        upper: The maximum value of the tensor. Shape (dims,)

    Returns:
        Denormalized transform of the tensor. Shape (N, dims)
    """
    # default value of center
    offset = (lower + upper) * 0.5
    # return normalized tensor
    return x * (upper - lower) * 0.5 + offset

@torch.jit.script
def saturate(x: torch.Tensor, lower: torch.Tensor, upper: torch.Tensor) -> torch.Tensor:
    """
    Clamps a given input tensor to (lower, upper).

    @note It uses pytorch broadcasting functionality to deal with batched input.

    Args:
        x: Input tensor of shape (N, dims).
        lower: The minimum value of the tensor. Shape (dims,)
        upper: The maximum value of the tensor. Shape (dims,)

    Returns:
        Clamped transform of the tensor. Shape (N, dims)
    """
    return torch.max(torch.min(x, upper), lower)

"""
Rotation conversions
"""

@torch.jit.script
def quat_diff_rad(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Get the difference in radians between two quaternions.

    Args:
        a: first quaternion, shape (N, 4)
        b: second quaternion, shape (N, 4)
    Returns:
        Difference in radians, shape (N,)
    """
    b_conj = quat_conjugate(b)
    mul = quat_mul(a, b_conj)
    # 2 * torch.acos(torch.abs(mul[:, -1]))
    return 2.0 * torch.asin(
        torch.clamp(
            torch.norm(
                mul[:, 0:3],
                p=2, dim=-1), max=1.0)
    )


@torch.jit.script
def local_to_world_space(pos_offset_local: torch.Tensor, pose_global: torch.Tensor):
    """ Convert a point from the local frame to the global frame
    Args:
        pos_offset_local: Point in local frame. Shape: [N, 3]
        pose_global: The spatial pose of this point. Shape: [N, 7]
    Returns:
        Position in the global frame. Shape: [N, 3]
    """
    quat_pos_local = torch.cat(
        [pos_offset_local, torch.zeros(pos_offset_local.shape[0], 1, dtype=torch.float32, device=pos_offset_local.device)],
        dim=-1
    )
    quat_global = pose_global[:, 3:7]
    quat_global_conj = quat_conjugate(quat_global)
    pos_offset_global = quat_mul(quat_global, quat_mul(quat_pos_local, quat_global_conj))[:, 0:3]

    result_pos_gloal = pos_offset_global + pose_global[:, 0:3]

    return result_pos_gloal

# NB: do not make this function jit, since it is passed around as an argument.
def normalise_quat_in_pose(pose):
    """Takes a pose and normalises the quaternion portion of it.

    Args:
        pose: shape N, 7
    Returns:
        Pose with normalised quat. Shape N, 7
    """
    pos = pose[:, 0:3]
    quat = pose[:, 3:7]
    quat /= torch.norm(quat, dim=-1, p=2).reshape(-1, 1)
    return torch.cat([pos, quat], dim=-1)

@torch.jit.script
def my_quat_rotate(q, v):
    shape = q.shape
    q_w = q[:, -1]
    q_vec = q[:, :3]
    a = v * (2.0 * q_w ** 2 - 1.0).unsqueeze(-1)
    b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
    c = q_vec * \
        torch.bmm(q_vec.view(shape[0], 1, 3), v.view(
            shape[0], 3, 1)).squeeze(-1) * 2.0
    return a + b + c

@torch.jit.script
def quat_to_angle_axis(q):
    # type: (Tensor) -> Tuple[Tensor, Tensor]
    # computes axis-angle representation from quaternion q
    # q must be normalized
    min_theta = 1e-5
    qx, qy, qz, qw = 0, 1, 2, 3

    sin_theta = torch.sqrt(1 - q[..., qw] * q[..., qw])
    angle = 2 * torch.acos(q[..., qw])
    angle = normalize_angle(angle)
    sin_theta_expand = sin_theta.unsqueeze(-1)
    axis = q[..., qx:qw] / sin_theta_expand

    mask = sin_theta > min_theta
    default_axis = torch.zeros_like(axis)
    default_axis[..., -1] = 1

    angle = torch.where(mask, angle, torch.zeros_like(angle))
    mask_expand = mask.unsqueeze(-1)
    axis = torch.where(mask_expand, axis, default_axis)
    return angle, axis

@torch.jit.script
def angle_axis_to_exp_map(angle, axis):
    # type: (Tensor, Tensor) -> Tensor
    # compute exponential map from axis-angle
    angle_expand = angle.unsqueeze(-1)
    exp_map = angle_expand * axis
    return exp_map

@torch.jit.script
def quat_to_exp_map(q):
    # type: (Tensor) -> Tensor
    # compute exponential map from quaternion
    # q must be normalized
    angle, axis = quat_to_angle_axis(q)
    exp_map = angle_axis_to_exp_map(angle, axis)
    return exp_map

def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to rotation matrices.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1)
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    mat = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    return mat.reshape(quaternions.shape[:-1] + (3, 3))

def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
    """
    Returns torch.sqrt(torch.max(0, x))
    subgradient is zero where x is 0.
    """
    ret = torch.zeros_like(x)
    positive_mask = x > 0
    ret[positive_mask] = torch.sqrt(x[positive_mask])
    return ret


def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as rotation matrices to quaternions.
    Args:
        matrix: Rotation matrices as tensor of shape (..., 3, 3).
    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")

    batch_dim = matrix.shape[:-2]
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
        matrix.reshape(batch_dim + (9,)), dim=-1
    )

    q_abs = _sqrt_positive_part(
        torch.stack(
            [
                1.0 + m00 + m11 + m22,
                1.0 + m00 - m11 - m22,
                1.0 - m00 + m11 - m22,
                1.0 - m00 - m11 + m22,
            ],
            dim=-1,
        )
    )

    quat_by_rijk = torch.stack(
        [
            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
        ],
        dim=-2,
    )

    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))

    return quat_candidates[
        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :
    ].reshape(batch_dim + (4,))


@torch.jit.script
def quat_to_tan_norm(q):
    # type: (Tensor) -> Tensor
    # represents a rotation using the tangent and normal vectors
    ref_tan = torch.zeros_like(q[..., 0:3])
    ref_tan[..., 0] = 1
    tan = my_quat_rotate(q, ref_tan)
    
    ref_norm = torch.zeros_like(q[..., 0:3])
    ref_norm[..., -1] = 1
    norm = my_quat_rotate(q, ref_norm)
    
    norm_tan = torch.cat([tan, norm], dim=len(tan.shape) - 1)
    return norm_tan

@torch.jit.script
def euler_xyz_to_exp_map(roll, pitch, yaw):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    q = quat_from_euler_xyz(roll, pitch, yaw)
    exp_map = quat_to_exp_map(q)
    return exp_map

@torch.jit.script
def exp_map_to_angle_axis(exp_map):
    min_theta = 1e-5

    angle = torch.norm(exp_map, p=2, dim=-1)
    angle_exp = torch.unsqueeze(angle, dim=-1)
    axis = exp_map / angle_exp
    angle = normalize_angle(angle)

    default_axis = torch.zeros_like(exp_map)
    default_axis[..., -1] = 1

    mask = angle > min_theta
    angle = torch.where(mask, angle, torch.zeros_like(angle))
    mask_expand = mask.unsqueeze(-1)
    axis = torch.where(mask_expand, axis, default_axis)

    return angle, axis

@torch.jit.script
def exp_map_to_quat(exp_map):
    angle, axis = exp_map_to_angle_axis(exp_map)
    q = quat_from_angle_axis(angle, axis)
    return q

@torch.jit.script
def slerp(q0, q1, t):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    qx, qy, qz, qw = 0, 1, 2, 3

    cos_half_theta = q0[..., qw] * q1[..., qw] \
                   + q0[..., qx] * q1[..., qx] \
                   + q0[..., qy] * q1[..., qy] \
                   + q0[..., qz] * q1[..., qz]
    
    neg_mask = cos_half_theta < 0
    q1 = q1.clone()
    q1[neg_mask] = -q1[neg_mask]
    cos_half_theta = torch.abs(cos_half_theta)
    cos_half_theta = torch.unsqueeze(cos_half_theta, dim=-1)

    half_theta = torch.acos(cos_half_theta);
    sin_half_theta = torch.sqrt(1.0 - cos_half_theta * cos_half_theta)

    ratioA = torch.sin((1 - t) * half_theta) / sin_half_theta
    ratioB = torch.sin(t * half_theta) / sin_half_theta; 
    
    new_q_x = ratioA * q0[..., qx:qx+1] + ratioB * q1[..., qx:qx+1]
    new_q_y = ratioA * q0[..., qy:qy+1] + ratioB * q1[..., qy:qy+1]
    new_q_z = ratioA * q0[..., qz:qz+1] + ratioB * q1[..., qz:qz+1]
    new_q_w = ratioA * q0[..., qw:qw+1] + ratioB * q1[..., qw:qw+1]

    cat_dim = len(new_q_w.shape) - 1
    new_q = torch.cat([new_q_x, new_q_y, new_q_z, new_q_w], dim=cat_dim)

    new_q = torch.where(torch.abs(sin_half_theta) < 0.001, 0.5 * q0 + 0.5 * q1, new_q)
    new_q = torch.where(torch.abs(cos_half_theta) >= 1, q0, new_q)

    return new_q

@torch.jit.script
def calc_heading(q):
    # type: (Tensor) -> Tensor
    # calculate heading direction from quaternion
    # the heading is the direction on the xy plane
    # q must be normalized
    ref_dir = torch.zeros_like(q[..., 0:3])
    ref_dir[..., 0] = 1
    rot_dir = my_quat_rotate(q, ref_dir)

    heading = torch.atan2(rot_dir[..., 1], rot_dir[..., 0])
    return heading

@torch.jit.script
def calc_heading_quat(q):
    # type: (Tensor) -> Tensor
    # calculate heading rotation from quaternion
    # the heading is the direction on the xy plane
    # q must be normalized
    heading = calc_heading(q)
    axis = torch.zeros_like(q[..., 0:3])
    axis[..., 2] = 1

    heading_q = quat_from_angle_axis(heading, axis)
    return heading_q

@torch.jit.script
def calc_heading_quat_inv(q):
    # type: (Tensor) -> Tensor
    # calculate heading rotation from quaternion
    # the heading is the direction on the xy plane
    # q must be normalized
    heading = calc_heading(q)
    axis = torch.zeros_like(q[..., 0:3])
    axis[..., 2] = 1

    heading_q = quat_from_angle_axis(-heading, axis)
    return heading_q


# EOF

NCF_policies/isaacgymenvs/utils/torch_jit_utils.py
carolinahiguera/NCF-Policy
# Copyright (c) 2018-2023, NVIDIA Corporation
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# python
#import pwd
import getpass
import tempfile
import time
from collections import OrderedDict
from os.path import join

import numpy as np
import torch
import random
import os


def retry(times, exceptions):
    """
    Retry Decorator https://stackoverflow.com/a/64030200/1645784
    Retries the wrapped function/method `times` times if the exceptions listed
    in ``exceptions`` are thrown
    :param times: The number of times to repeat the wrapped function/method
    :type times: Int
    :param exceptions: Lists of exceptions that trigger a retry attempt
    :type exceptions: Tuple of Exceptions
    """
    def decorator(func):
        def newfn(*args, **kwargs):
            attempt = 0
            while attempt < times:
                try:
                    return func(*args, **kwargs)
                except exceptions:
                    print(f'Exception thrown when attempting to run {func}, attempt {attempt} out of {times}')
                    time.sleep(min(2 ** attempt, 30))
                    attempt += 1

            return func(*args, **kwargs)
        return newfn
    return decorator


def flatten_dict(d, prefix='', separator='.'):
    res = dict()
    for key, value in d.items():
        if isinstance(value, (dict, OrderedDict)):
            res.update(flatten_dict(value, prefix + key + separator, separator))
        else:
            res[prefix + key] = value

    return res


def set_np_formatting():
    """ formats numpy print """
    np.set_printoptions(edgeitems=30, infstr='inf',
                        linewidth=4000, nanstr='nan', precision=2,
                        suppress=False, threshold=10000, formatter=None)


def set_seed(seed, torch_deterministic=False, rank=0):
    """ set seed across modules """
    if seed == -1 and torch_deterministic:
        seed = 42 + rank
    elif seed == -1:
        seed = np.random.randint(0, 10000)
    else:
        seed = seed + rank

    print("Setting seed: {}".format(seed))

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if torch_deterministic:
        # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.use_deterministic_algorithms(True)
    else:
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False

    return seed

def nested_dict_set_attr(d, key, val):
    pre, _, post = key.partition('.')
    if post:
        nested_dict_set_attr(d[pre], post, val)
    else:
        d[key] = val
    
def nested_dict_get_attr(d, key):
    pre, _, post = key.partition('.')
    if post:
        return nested_dict_get_attr(d[pre], post)
    else:
        return d[key]

def ensure_dir_exists(path):
    if not os.path.exists(path):
        os.makedirs(path)
    return path


def safe_ensure_dir_exists(path):
    """Should be safer in multi-treaded environment."""
    try:
        return ensure_dir_exists(path)
    except FileExistsError:
        return path


def get_username():
    uid = os.getuid()
    try:
        return getpass.getuser()
    except KeyError:
        # worst case scenario - let's just use uid
        return str(uid)


def project_tmp_dir():
    tmp_dir_name = f'ige_{get_username()}'
    return safe_ensure_dir_exists(join(tempfile.gettempdir(), tmp_dir_name))

# EOF

NCF_policies/isaacgymenvs/utils/utils.py
carolinahiguera/NCF-Policy
"""Installation script for the 'isaacgymenvs' python package."""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

from setuptools import setup, find_packages

import os

root_dir = os.path.dirname(os.path.realpath(__file__))


# Minimum dependencies required prior to installation
INSTALL_REQUIRES = [
    # RL
    "gym==0.23.1",
    "torch",
    "omegaconf",
    "termcolor",
    "jinja2",
    "hydra-core>=1.1",
    "rl-games>=1.6.0",
    "pyvirtualdisplay",
    ]



# Installation operation
setup(
    name="ncf_envs",
    author="",
    version="1.4.0",
    description="Benchmark environments for testing NCF on downstream tasks.",
    keywords=["robotics", "ncf", "rl"],
    include_package_data=True,
    python_requires=">=3.6",
    install_requires=INSTALL_REQUIRES,
    packages=find_packages("."),
    classifiers=["Natural Language :: English", "Programming Language :: Python :: 3.6, 3.7, 3.8"],
    zip_safe=False,
)

# EOF

NCF_policies/setup.py
carolinahiguera/NCF-Policy