Ensemble-of-Costs Guided Diffusion For Motion Planning
import pickle
import torch
import sys
import numpy as np
import os

current_dir = os.getcwd()

sys.path.append(f'{current_dir}/motion-policy-networks/') # path to mpinets

from mpinets.types import PlanningProblem, ProblemSet
from geometrout.primitive import Cuboid, Cylinder
from robofin.robots import FrankaRobot

class TestDataset:
    def __init__(self, type='global', d_path='./datasets/') -> None:
        '''

        Parameters:
        type: 'global' or 'hybrid' or 'both'
        d_path: path to the datasets directory
        '''
        self.type = type

        if d_path[-1]!='/':
            d_path += '/'

        if type=='global':
            # self.data = np.load(d_path + 'global_solvable_problems.pkl', allow_pickle=True)
            with open(d_path + f'global_solvable_problems.pkl', 'rb') as f:
                self.data = pickle.load(f)
        elif type=='hybrid':
            with open(d_path + f'hybrid_solvable_problems.pkl', 'rb') as f:
                self.data = pickle.load(f)
        elif type=='both':
            with open(d_path + f'both_solvable_problems.pkl', 'rb') as f:
                self.data = pickle.load(f)
        else:
            raise ModuleNotFoundError("No such dataset exists")
    
        print('-'*50)
        print(f"Loaded the dataset with {self.data.keys()} scene types, each consisting of {self.data['tabletop'].keys()} problem types")
        print('-'*50)
        print("Dataset statistics:"+'*'*20)
        total_data = 0
        for key1 in self.data.keys():
            for key2 in self.data[key1]:
                print(f"{key1}: {key2}: {len(self.data[key1][key2])} problems")
                total_data += len(self.data[key1][key2])
        print('-'*50)
        print(f"Total data: {total_data}")
        print('-'*50)

        self.tabletop_data = np.hstack((self.data['tabletop']['task_oriented'], self.data['tabletop']['neutral_start'], self.data['tabletop']['neutral_goal']))
        self.cubby_data = np.hstack((self.data['cubby']['task_oriented'], self.data['cubby']['neutral_start'], self.data['cubby']['neutral_goal']))
        self.merged_cubby_data = np.hstack((self.data['merged_cubby']['task_oriented'], self.data['merged_cubby']['neutral_start'], self.data['merged_cubby']['neutral_goal']))
        self.dresser_data = np.hstack((self.data['dresser']['task_oriented'], self.data['dresser']['neutral_start'], self.data['dresser']['neutral_goal']))

        self.data_nums = {}
        self.data_nums['tabletop'] = len(self.tabletop_data)
        self.data_nums['cubby'] = len(self.cubby_data)
        self.data_nums['merged_cubby'] = len(self.cubby_data)
        self.data_nums['dresser'] = len(self.dresser_data)


    def fetch_batch(self, problem_type='random', task_type='random'):
        '''
        Parameters:
        scene_type: 'all' or 'tabletop' or 'cubby' or 'merged_cubby' or 'dresser'
        problem_type: 'random' or 'task_oriented' or 'neutral_start' or 'neutral_goal'
        '''
        raise NotImplementedError("Not implemented yet!")
    
    

    def fetch_data(self, scene_num, scene_type='tabletop'):
        '''
        Parameters:
        scene_type: tabletop or cubby or merged_cubby or dresser
        '''
        if scene_type=='tabletop':
            data = self.tabletop_data[scene_num]
        elif scene_type=='cubby':
            data = self.cubby_data[scene_num]
        elif scene_type=='merged_cubby':
            data = self.merged_cubby_data[scene_num]
        elif scene_type=='dresser':
            data = self.dresser_data[scene_num]
        else:
            raise ModuleNotFoundError("What are you looking for? This dataset only has 4 options. Try to choose one of them and retry!")

        # Initialize lists to store centers and quaternions
        cuboid_centers = []
        cuboid_dims = []
        cuboid_quaternions = []

        cylinder_centers = []
        cylinder_heights = []
        cylinder_radii = []
        cylinder_quaternions = []

        num_cuboids = 0
        num_cylinders = 0

        # Extract centers and quaternions based on object type
        for obstacle in data.obstacles:
            if isinstance(obstacle, Cuboid):
                cuboid_centers.append(np.array(obstacle.center))
                cuboid_quaternions.append(np.array(list(obstacle._pose._so3._quat)))
                cuboid_dims.append(np.array(obstacle.dims))
                num_cuboids += 1
            elif isinstance(obstacle, Cylinder):
                cylinder_centers.append(np.array(obstacle.center))
                cylinder_heights.append(np.array(obstacle.height))
                cylinder_radii.append(np.array(obstacle.radius))
                cylinder_quaternions.append(np.array(list(obstacle._pose._so3._quat)))
                num_cylinders += 1

        # Convert the lists to NumPy arrays
        cuboid_config = []
        cylinder_config = []
        if(num_cuboids >= 1):
            # Convert the lists to NumPy arrays
            cuboid_centers = np.array(cuboid_centers)
            cuboid_dims = np.array(cuboid_dims)
            cuboid_quaternions = np.roll(np.array(cuboid_quaternions), -1, axis = 1)
            cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1)

        if(num_cylinders >= 1):
            cylinder_centers = np.array(cylinder_centers)
            cylinder_heights = np.array(cylinder_heights).reshape(-1, 1)
            cylinder_radii = np.array(cylinder_radii).reshape(-1, 1)
            cylinder_quaternions = np.roll(np.array(cylinder_quaternions), -1, axis = 1)
            cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1)

            cylinder_cuboid_dims = np.zeros((num_cylinders, 3))
            cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0]

        if(num_cylinders >= 1):
            obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0)
            obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0)
            obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0)
        else: 
            obstacle_centers = cuboid_centers
            obstacle_dims = cuboid_dims
            obstacle_quaternions = cuboid_quaternions

        obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1)
        # cuboid_centers = np.array(cuboid_centers)
        # cuboid_dims = np.array(cuboid_dims)
        # cuboid_quaternions = np.array(cuboid_quaternions)
        # cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1)

        # cylinder_centers = np.array(cylinder_centers)
        # cylinder_heights = np.array(cylinder_heights).reshape(-1, 1)
        # cylinder_radii = np.array(cylinder_radii).reshape(-1, 1)
        # cylinder_quaternions = np.array(cylinder_quaternions)
        # cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1)

        # cylinder_cuboid_dims = np.zeros((num_cylinders, 3))
        # cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0]
        # cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0]
        # cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0]

        # obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0)
        # obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0)
        # obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0)

        # obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1)

        start = data.q0
        goal_ee = data.target

        joint_7_samples = np.concatenate((np.random.uniform(-2.8973, 2.8973, 50), np.linspace(-2.8973, 2.8973, 50)))
        all_ik_goals = []
        for ik_sol_num, joint_ang in enumerate(joint_7_samples):
            ik_solutions = np.array(FrankaRobot.ik(goal_ee, joint_ang, 'right_gripper')).reshape((-1, 7))
            if len(ik_solutions)==0:
                continue
        
            if len(all_ik_goals)==0:
                all_ik_goals = ik_solutions.reshape((-1, 7))
            else:
                all_ik_goals = np.vstack((all_ik_goals, ik_solutions[0].reshape((-1, 7)))) # n*7
            
        # print(f"Total IK Solutions: {len(all_ik_goals)}")
        return obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders, start, all_ik_goals
    
    def random(self):
        # Initialize lists to store centers and quaternions
        cuboid_centers = []
        cuboid_dims = []
        cuboid_quaternions = []

        cylinder_centers = []
        cylinder_heights = []
        cylinder_radii = []
        cylinder_quaternions = []

        num_cuboids = 0
        num_cylinders = 0

        cuboid_config, cylinder_config = None, None

        # Extract centers and quaternions based on object type
        for obstacle in obstacles:
            if isinstance(obstacle, Cuboid):
                cuboid_centers.append(np.array(obstacle.center))
                cuboid_quaternions.append(np.array(list(obstacle._pose._so3._quat)))
                cuboid_dims.append(np.array(obstacle.dims))
                num_cuboids += 1
            elif isinstance(obstacle, Cylinder):
                cylinder_centers.append(np.array(obstacle.center))
                cylinder_heights.append(np.array(obstacle.height))
                cylinder_radii.append(np.array(obstacle.radius))
                cylinder_quaternions.append(np.array(list(obstacle._pose._so3._quat)))
                num_cylinders += 1

        if(num_cuboids >= 1):
            # Convert the lists to NumPy arrays
            cuboid_centers = np.array(cuboid_centers)
            cuboid_dims = np.array(cuboid_dims)
            cuboid_quaternions = np.roll(np.array(cuboid_quaternions), -1, axis = 1)
            cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1)

        if(num_cylinders >= 1):
            cylinder_centers = np.array(cylinder_centers)
            cylinder_heights = np.array(cylinder_heights).reshape(-1, 1)
            cylinder_radii = np.array(cylinder_radii).reshape(-1, 1)
            cylinder_quaternions = np.roll(np.array(cylinder_quaternions),  -1, axis = 1)
            # print('cyl centres shape: {}'.format(cylinder_centers.shape))
            cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1)

            cylinder_cuboid_dims = np.zeros((num_cylinders, 3))
            cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0]

        if(num_cylinders >= 1):
            obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0)
            obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0)
            obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0)
        else: 
            obstacle_centers = cuboid_centers
            obstacle_dims = cuboid_dims
            obstacle_quaternions = cuboid_quaternions

        obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1)

        return obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders
        
        
if __name__=='__main__':
    tData = TestDataset()

    tData.fetch_data(scene_num=5)

datasets/load_test_dataset.py
vishal-2000/EDMP
from .models import *
from .diffusion import *
from .gaussian import *
diffusion/__init__.py
vishal-2000/EDMP
import numpy as np
import torch
import copy
import time

from .gaussian import Gaussian

class Diffusion(Gaussian):
    
    def __init__(self, T, device, variance_thresh = 0.02):

        self.T = T
        self.beta = self.schedule_variance(variance_thresh)

        self.alpha = 1 - self.beta
        self.alpha_bar = np.array(list(map(lambda t:np.prod(self.alpha[:t]), np.arange(self.T+1)[1:])))

        self.device = device

        Gaussian.__init__(self)
    
    def cosine_func(self, t):
        """
        Implements the cosine function at a timestep t
        Inputs:
        t -> Integer ; current timestep
        T -> Integer ; total number of timesteps
        Outputs:
        out -> Float ; cosine function output
        """
        
        s = 1e-10
        out = np.cos((t/self.T + s) * (np.pi/2) / (1 + s)) ** 0.15
        
        return out

    def schedule_variance(self, thresh = 0.02):
        """
        Schedules the variance for the diffuser
        Inputs:
        T      -> Integer ; total number of timesteps
        thresh -> Float ; variance threshold at the last step
        Outputs:
        schedule -> Numpy array of shape (2,) ; the variance schedule
        """
            
        schedule = np.linspace(0, thresh, self.T + 1)[1:]

        return schedule

    def q_sample(self, x, t, eps = None):
        """
        Generates q(xt+1/xt)
        Inputs:
        x0 -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        t -> Numpy array of length (num_samples, ); timesteps (these are the output timesteps)
        Outputs:
        xt -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        mean -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean from which sample is taken
        var-> Numpy array of length (num_samples); variance from which sample is taken
        """

        # Gather trajectory params:
        b = x.shape[0]
        c = x.shape[1]
        n = x.shape[2]
        
        if type(eps) == type(None):
            eps = np.random.multivariate_normal(mean = np.zeros((n*c)), 
                                                cov = np.eye(n*c),
                                                size = (b,)).reshape(b, c, n)
        
        xt = np.sqrt(self.alpha[t-1, np.newaxis, np.newaxis]) * x + np.sqrt(1 - self.alpha[t-1, np.newaxis, np.newaxis]) * eps
        mean = (np.sqrt(self.alpha[t-1, np.newaxis, np.newaxis]) * x)
        var = np.sqrt(1 - self.alpha[t-1])

        return xt, mean, var
    
    def q_sample_from_x0(self, x0, t, eps = None):
        """
        Generates q(xt/x0)
        Inputs:
        x0 -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        t -> Numpy array of length (num_samples, ); timesteps (these are the output timesteps)
        Outputs:
        xt -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        mean -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean from which sample is taken
        var-> Numpy array of length (num_samples); variance from which sample is taken
        """

        # Gather trajectory params:
        b = x0.shape[0]
        c = x0.shape[1]
        n = x0.shape[2]
        
        if type(eps) == type(None):
            eps = np.random.multivariate_normal(mean = np.zeros((n*c)), 
                                                cov = np.eye(n*c),
                                                size = (b,)).reshape(b, c, n)
        
        xt = np.sqrt(self.alpha_bar[t-1, np.newaxis, np.newaxis]) * x0 + np.sqrt(1 - self.alpha_bar[t-1, np.newaxis, np.newaxis]) * eps
        mean = (np.sqrt(self.alpha_bar[t-1, np.newaxis, np.newaxis]) * x0)
        var = np.sqrt(1 - self.alpha_bar[t-1, np.newaxis, np.newaxis])

        return xt, mean, var
    
    def p_sample(self, xt, t, eps):
        """
        Generates reverse probability p(xt-1/xt, eps) with no dependence on x0.
        """

        xt_prev = (xt - np.sqrt(1 - self.alpha[t-1]) * eps) / np.sqrt(self.alpha[t-1])

        return xt_prev
    
    def p_sample_using_posterior(self, xt, t, eps):
        """
        Generates reverse probability p(xt-1/xt, x0) using posterior mean and posterior variance.
        """        
        
        # Gather trajectory params:
        b = xt.shape[0]
        c = xt.shape[1]
        n = xt.shape[2]

        z = np.random.multivariate_normal(mean = np.zeros(n), cov = np.eye(n), size = (b, c))
        z[np.where(t == 1), :, :] = 0

        alpha = self.alpha[t-1, np.newaxis, np.newaxis]
        alpha_bar = self.alpha_bar[t-1, np.newaxis, np.newaxis]
        beta = self.beta[t-1, np.newaxis, np.newaxis]

        xt_prev = (xt - ((1 - alpha) / np.sqrt(1 - alpha_bar)) * eps) / np.sqrt(alpha) + beta*z

        return xt_prev.copy()

    def forward_diffuse(self, x0, condition = True):
        """
        Forward diffuses the trajectory till the last timestep T
        Inputs:
        x -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        Outputs:
        diffusion_trajs -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); diffused trajectories
        eps -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); Noise added at each timestep
        kl_divs -> Numpy array of shape (Timesteps, num_samples); KL Divergence at each timestep
        """
                
        # Gather trajectory params:
        b = x0.shape[0]
        c = x0.shape[1]
        n = x0.shape[2]
        
        # Initialize the diffusion trajectories across time steps:
        diffusion_trajs = np.zeros((self.T + 1, b, c, n))
        diffusion_trajs[0] = x0.copy()

        # Calculate epsilons to forward diffuse current trajectory:
        mean = np.zeros(n*c)
        cov = np.eye(n*c)
        eps = np.random.multivariate_normal(mean, cov, size = (self.T, b))
        kl_divs = np.zeros((self.T, b))

        # ??? This loop can probably be vectorized
        
        for t in range(1, self.T + 1):

            diffusion_trajs[t] = self.q_sample(diffusion_trajs[t-1], t, eps[t-1])
            if condition:
                diffusion_trajs[t, :, :, 0] = x0[:, :, 0].copy()
                diffusion_trajs[t, :, :, -1] = x0[:, :, -1].copy()
            
            # ???
            kl_divs[t-1, :] = self.KL_divergence_against_gaussian(diffusion_trajs[t].flatten())

        return diffusion_trajs.copy(), eps.copy(), kl_divs.copy()

    def reverse_diffuse(self, xT, eps):
        """
        Reverse diffuses the trajectory from last timestep T to first timestep 0
        Inputs:
        xT  -> Numpy array of shape (num_samples, num_channels, trajectory_length)
        eps -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); Noise added at each timestep
        Outputs:
        diffusion_trajs -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); reverse diffused trajectories
        """
        
        # Gather trajectory params:
        b = xT.shape[0]
        c = xT.shape[1]
        n = xT.shape[2]
        
        diffusion_trajs = np.zeros((self.T + 1, b, c, n))
        diffusion_trajs[self.T] = xT.copy()

        for t in range(self.T, 0, -1):

            diffusion_trajs[t-1] = self.p_sample(diffusion_trajs[t], t, eps[t-1])

        return diffusion_trajs.copy()

    def generate_q_sample(self, x0, time_steps = None, condition = True, return_type = "tensor"):
        """
        Generates q samples for a random set of timesteps, useful for training\n\n
        Inputs:\n
        x0          -> Numpy array of shape (num_samples, num_channels, trajectory_length)\n
        condition   -> Bool ; whether to apply conditioning or not\n
        return_type -> String (either "tensor" or "numpy") ; whether to return a tensor or numpy array\n\n
        Outputs:\n
        X -> Numpy array of shape (num_samples, num_channels, trajectory_length); xt from each x0 diffused to a random timestep t\n
        Y -> Numpy array of shape (num_samples, num_channels, trajectory_length); the noise added to each x0\n
        time_step -> Numpy array of shape (num_samples, ); the timestep to which each x0 is diffused\n
        means -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean of each x0 diffused to a random timestep t\n
        vars -> Numpy array of shape (num_samples, ); the variance at the timestep to which x0 is diffused\n
        """

        # Refer to the Training Algorithm in Ho et al 2020 for the psuedo code of this function
            
        # Gather trajectory params:
        b = x0.shape[0]
        c = x0.shape[1]
        n = x0.shape[2]
        
        if type(time_steps) == type(None):
            time_steps = np.random.randint(1, self.T + 1, size = (b, ))

        # Remember, for each channel, we get a multivariate normal distribution.
        eps = np.random.multivariate_normal(mean = np.zeros((n,)), cov = np.eye(n), size = (b, c))
        
        # Size chart:
        # x0         => (num_samples, 2, traj_len)
        # xt         => (num_samples, 2, traj_len)
        # alpha_bar  => (num_samples, 1, 1)
        # eps        => (num_samples, 2, traj_len)
        
        xt, means, vars = self.q_sample_from_x0(x0, time_steps, eps)  

        # CONDITIONING:
        if condition:
            xt[:, :, 0] = x0[:, :, 0].copy()
            xt[:, :, -1] = x0[:, :, -1].copy()

        # Type Conversion:
        if return_type == "tensor":
            X = torch.tensor(xt, dtype = torch.float32)
            Y = torch.tensor(eps, dtype = torch.float32)
            time_steps = torch.tensor(time_steps, dtype = torch.float32)
        elif return_type == "numpy":
            X = xt.copy()
            Y = eps.copy()

        return X, Y, time_steps, means, vars
    
    def denoise(self, model, traj_len, num_channels, start = None, goal = None, condition = True):
        
        X_t = np.random.multivariate_normal(mean = np.zeros(traj_len), cov = np.eye(traj_len), size = (1, num_channels))

        if condition:
            X_t[:, :, 0] = start[:]
            X_t[:, :, -1] = goal[:]

        model.train(False)

        for t in range(self.T, 0, -1):

            X_input = torch.tensor(X_t, dtype = torch.float32).to(self.device)
            time_in = torch.tensor([t], dtype = torch.float32).to(self.device)

            epsilon = model(X_input, time_in).numpy(force=True)

            X_t = self.p_sample_using_posterior(X_t, t, epsilon)

            if condition:
                X_t[:, :, 0] = start[:]
                X_t[:, :, -1] = goal[:]

            print("\rDenoised " + str(t-1) + " timesteps", end="", flush=True)

        return X_t[0]
    
    def clip_joints(self, joints):

        joint_lower_limits = np.array([-166*(np.pi/180), 
                                           -101*(np.pi/180), 
                                           -166*(np.pi/180), 
                                           -176*(np.pi/180),
                                           -166*(np.pi/180),
                                           -1*(np.pi/180),
                                           -166*(np.pi/180)])
        
        joint_upper_limits = np.array([166*(np.pi/180), 
                                           101*(np.pi/180), 
                                           166*(np.pi/180), 
                                           -4*(np.pi/180),
                                           166*(np.pi/180),
                                           215*(np.pi/180),
                                           166*(np.pi/180)])

        return np.clip(joints, joint_lower_limits[np.newaxis, :, np.newaxis], joint_upper_limits[np.newaxis, :, np.newaxis])
    
    def denoise_guided(self, model, guide, traj_len, num_channels, batch_size = 1, start = None, goal = None, 
                       condition = True, benchmarking = False, guidance_schedule='constant', guidance_scale=0.05):
        
        X_t = np.random.multivariate_normal(mean = np.zeros(traj_len), cov = np.eye(traj_len), size = (batch_size, num_channels))

        if condition:
            X_t[:, :, 0] = start[:]
            X_t[:, :, -1] = goal[:]

        model.train(False) 

        period = 2

        # full_info = np.zeros((self.T+1, batch_size, 7, 50))
        for t in range(self.T, 0, -1):
            print(f"\rDenoising: {t}", end="")

            # full_info[t, :, : ,:] = X_t.copy()  # copy.deepcopy(X_t)

            X_input = torch.tensor(X_t, dtype = torch.float32).to(self.device)
            time_in = torch.tensor([t], dtype = torch.float32).to(self.device)

            epsilon = model(X_input, time_in).numpy(force=True)

            X_t = self.p_sample_using_posterior(X_t, t, epsilon)

            if (t%period) < (period/2):
                if t >= 5:
                    clipped_joints = self.clip_joints(X_t[:, :, 1:-1])
                    # st = time.time()
                    gradient = guide.get_gradient(clipped_joints, start[:], goal[:], t)
                    # print(f"\nGradient time: {time.time() - st}")
                    if guidance_schedule == 'constant':
                        X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - guidance_scale * gradient
                    elif guidance_schedule == 'varying':
                        X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - (1.4 + (t/self.T)) * gradient
                    else:
                        raise NotImplementedError("This type of guidance schedule is not implemented!!")
                    
                    # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - 1.4 * gradient

                # weight = np.clip(np.log(1+((t-2)/self.T)*(np.exp(1) - 1)), 0.005, 1) * 2.2
                # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - weight * gradient
                # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - 2. * gradient

            if condition:
                X_t[:, :, 0] = start[:]
                X_t[:, :, -1] = goal[:]

            if not benchmarking:
                print("\rDenoised " + str(t-1) + " timesteps", end="", flush=True)

        # full_info[0, :, :, :] = X_t.copy()

        return X_t.copy() #, full_info.copy()
        
    




diffusion/diffusion.py
vishal-2000/EDMP
import numpy as np
import matplotlib.pyplot as plt

class Gaussian:

    def pdf(self, x, mu, sigma):
        """
        Returns the probability density function given some mu, sigma and a list of points x.
        """
        out = (1 / (sigma * np.sqrt(2*np.pi))) * np.exp(-0.5 * (((x - mu) / sigma) ** 2))

        return out

    def KL_divergence(self, distr1, distr2):
        """
        Computes the Kullback-Leibler Divergence between 2 distributions
        """

        return np.mean(distr1 * np.log(distr1 / distr2))

    def KL_divergence_against_gaussian(self, sample):
        """
        Computes the Kullback-Leibler Divergence of a given set of samples against a standard normal gaussian. Useful in analysing how
        the distribution shifts along the timesteps of diffusion.
        """

        mu = np.mean(sample)
        sigma = np.var(sample)

        x = np.linspace(-10, 10, int(1e+6))

        distr1 = self.pdf(x, 0, 1) 
        distr2 = self.pdf(x, mu, sigma)

        return self.KL_divergence(distr1, distr2)

    def get_limits(self, mu, sigma, edge_factor = 0.01):
        """
        Gets the upper and lower limit of the bell curve of the distribution, used for fixing limits for plotting the distribution.
        """

        p_min = edge_factor * self.pdf(mu, mu, sigma)
        x_1 = -1 * sigma * np.sqrt(2 * np.log(1 / (p_min * sigma * np.sqrt(2*np.pi)))) + mu
        x_2 = sigma * np.sqrt(2 * np.log(1 / (p_min * sigma * np.sqrt(2*np.pi)))) + mu

        return x_1, x_2

    def multivariate_pdf(self, mean, var, size = 1024, limits = (-1, 1)):
        """
        Inputs:
        mean => 1D array of size k, in a 2D trajectory it is k=2 and for 3D k=3
        covariance => 2D array of (k x k)
        size => Integer; Per axis length of output
        limits => Upper and lower limits for each axis
        Output:
        Array of 'k' dimensions of shape (size x size x .... x size)
        """

        mean = np.array(mean)
        k = mean.size

        axis_div = np.repeat([np.linspace(limits[0], limits[1], size)], repeats = k, axis = 0)
        x = np.array(np.meshgrid(*axis_div))

        mu = np.reshape(mean, (k, *np.ones(k, dtype = np.int32)))

        pdf = (1 / (((2 * np.pi) ** (k/2)) * np.sqrt(var**k))) * np.exp(-0.5 * np.sum((x - mu)**2, axis = 0) / var)

        return pdf


diffusion/gaussian.py
vishal-2000/EDMP
from .blocks import *
from .temporalunet import *
diffusion/models/__init__.py
vishal-2000/EDMP
import torch
import torch.nn as nn
from einops.layers.torch import Rearrange
import numpy as np
import einops

#################################################################################################################
#---------------------------------------------- LOWER LEVEL BLOCKS ---------------------------------------------#
#################################################################################################################

#--------------------------------------------- 1D CONVOLUTION BLOCK --------------------------------------------#

class Conv1dBlock(nn.Module):
    '''
        Conv1d --> GroupNorm --> Mish
    '''

    def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):
        
        super().__init__()

        self.block = nn.Sequential(
            nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2),
            Rearrange('batch channels horizon -> batch channels 1 horizon'),
            nn.GroupNorm(n_groups, out_channels),
            Rearrange('batch channels 1 horizon -> batch channels horizon'),
            nn.Mish(),
        )
        # Why is padding half of kernel size? This makes sure that in the first convolution, half the kernel has zero elements and the other
        # half has the first few elements of the input. The same applies to the last convolution.

    def forward(self, x):

        return self.block(x)
    
#----------------------------------------- SINUSOIDAL POSITION EMBEDDING ---------------------------------------#

class SinusoidalPosEmb(nn.Module):

    def __init__(self, dim, device):
        
        super().__init__()
        self.dim = dim
        self.device = device

    def forward(self, x):

        half_dim = self.dim // 2
        emb = np.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=self.device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)

        return emb

#------------------------------------------------- TEMPORAL MLP ------------------------------------------------#

class TimeMLP(nn.Module):

    def __init__(self, time_embed_dim, out_channels):

        super().__init__()
        
        self.time_mlp = nn.Sequential(
            nn.Mish(),
            nn.Linear(time_embed_dim, out_channels),
            Rearrange('batch t -> batch t 1'),      # This is probably extending by a dimension
        )

    def forward(self, t):

        return self.time_mlp(t)

#-------------------------------------------- INITIAL TIME EMBEDDING -------------------------------------------#

class TimeEmbedding(nn.Module):

    def __init__(self, dim, device):

        super().__init__()
        self.device = device
        
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(dim, self.device),
            nn.Linear(dim, dim * 4),
            nn.Mish(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, t):

        return self.time_mlp(t)

#################################################################################################################
#------------------------------------------ INTERMEDIATE LEVEL BLOCKS ------------------------------------------#
#################################################################################################################

#---------------------------------------------- LINEAR ATTENTION -----------------------------------------------#

class LinearAttention(nn.Module):
    
    def __init__(self, dim, heads=4, dim_head=32):
        
        super().__init__()
        
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads

        self.to_qkv = nn.Conv1d(dim, hidden_dim * 3, 1, bias=False)
        self.to_out = nn.Conv1d(hidden_dim, dim, 1)

    def forward(self, x):

        qkv = self.to_qkv(x).chunk(3, dim = 1)
        q, k, v = map(lambda t: einops.rearrange(t, 'b (h c) d -> b h c d', h=self.heads), qkv)

        q = q * self.scale
        k = k.softmax(dim = -1)

        # Produces the context of each element in k w.r.t all other elements in v. (weighted sum of every pair of rows possible)
        context = torch.einsum('b h d n, b h e n -> b h d e', k, v)

        # Weighted sum of every pair of columns possible:
        out = torch.einsum('b h d e, b h d n -> b h e n', context, q)
        
        # Recombine the the 3 chunks that where seperated to get back the same dimension as input
        out = einops.rearrange(out, 'b h c d -> b (h c) d')
        
        # Convolve back from hidden channels to the original number of channels (dim)
        out = self.to_out(out)

        return out

#------------------------------------------ RESIDUAL CONVOLUTION BLOCK ------------------------------------------#

class ResidualConvolutionBlock(nn.Module):

    def __init__(self, inp_channels, out_channels, time_embed_dim, kernel_size=5):
        super().__init__()

        self.blocks = nn.ModuleList([
            Conv1dBlock(inp_channels, out_channels, kernel_size),
            Conv1dBlock(out_channels, out_channels, kernel_size),
        ])

        self.time_mlp = TimeMLP(time_embed_dim, out_channels)

        # The convolution that forms the residual connection between input and output
        # If the input and the output have the same number of channels, this is an identity matrix.
        self.residual_conv = nn.Conv1d(inp_channels, out_channels, 1) \
            if inp_channels != out_channels else nn.Identity()

    def forward(self, x, t):
        '''
            x : [ batch_size x inp_channels x horizon ]
            t : [ batch_size x time_embed_dim ]
            returns:
            out : [ batch_size x out_channels x horizon ]
        '''
        
        out = self.blocks[0](x) + self.time_mlp(t)
        out = self.blocks[1](out)
        out = out + self.residual_conv(x)

        return out
    
#------------------------------------------- RESIDUAL ATTENTION BLOCK -------------------------------------------#

class ResidualAttentionBlock(nn.Module):

    def __init__(self, dim, eps = 1e-5):
        
        super().__init__()

        self.attention = LinearAttention(dim)

        # Layer Norm Parameters:
        self.eps = eps
        self.g = nn.Parameter(torch.ones(1, dim, 1))
        self.b = nn.Parameter(torch.zeros(1, dim, 1))

    def forward(self, x):

        # Layer Norm
        var = torch.var(x, dim=1, unbiased=False, keepdim=True)
        mean = torch.mean(x, dim=1, keepdim=True)
        out = (x - mean) / (var + self.eps).sqrt() * self.g + self.b

        # Attention Layer
        out = self.attention(out)

        # Residual Connection
        out = out + x

        return out

#################################################################################################################
#---------------------------------------------- HIGH LEVEL BLOCKS ----------------------------------------------#
#################################################################################################################

class DownSampler(nn.Module):
    
    def __init__(self, dim_in, dim_out, time_dim, is_last = False):

        super().__init__()
        
        self.down = nn.ModuleList([ResidualConvolutionBlock(dim_in, dim_out, time_embed_dim = time_dim),
                                   ResidualConvolutionBlock(dim_out, dim_out, time_embed_dim = time_dim),
                                   nn.Identity(),  # Replace this with ResidualAttentionBlock
                                   nn.Conv1d(dim_out, dim_out, kernel_size = 3, stride = 2, padding = 1) if not is_last else nn.Identity()])

    def forward(self, x, t):
        
        x = self.down[0](x, t)
        x = self.down[1](x, t)
        h = self.down[2](x)   
        out = self.down[3](h)

        return out, h

class MiddleBlock(nn.Module):
    
    def __init__(self, mid_dim, time_dim):

        super().__init__()

        self.middle = nn.ModuleList([ResidualConvolutionBlock(mid_dim, mid_dim, time_embed_dim = time_dim),
                                     nn.Identity(),  # Replace this with ResidualAttentionBlock
                                     ResidualConvolutionBlock(mid_dim, mid_dim, time_embed_dim = time_dim)])

    def forward(self, x, t):

        x = self.middle[0](x, t)
        x = self.middle[1](x)
        out = self.middle[2](x, t)

        return out

class UpSampler(nn.Module):
    
    def __init__(self, dim_in, dim_out, time_dim, is_last = False):

        super().__init__()

        self.up = nn.ModuleList([ResidualConvolutionBlock(dim_out * 2, dim_in, time_embed_dim = time_dim),
                                 ResidualConvolutionBlock(dim_in, dim_in, time_embed_dim = time_dim),
                                 nn.Identity(),  # Replace this with ResidualAttentionBlock
                                 nn.ConvTranspose1d(dim_in, dim_in, kernel_size = 4, stride = 2, padding = 1) if not is_last else nn.Identity()])

    def forward(self, x, h, t):

        x = torch.cat([x, h], dim = 1)

        x = self.up[0](x, t)
        x = self.up[1](x, t)
        x = self.up[2](x)   
        out = self.up[3](x)

        return out
diffusion/models/blocks.py
vishal-2000/EDMP
import torch
import torch.nn as nn
import numpy as np
import os
import torchvision.transforms.functional as tvtf

from diffusion.models.blocks import *

class TemporalUNet(nn.Module):

    def __init__(self, model_name, input_dim, time_dim, device, dims = (32, 64, 128, 256)):

        super(TemporalUNet, self).__init__()

        dims = [input_dim, *dims]  # length of dims is 5

        # Initial Time Embedding:
        self.time_embedding = TimeEmbedding(time_dim, device)

        # Down Sampling:
        self.down_samplers = nn.ModuleList([])
        for i in range(len(dims) - 2):      # Loops 0, 1, 2
            self.down_samplers.append(DownSampler(dims[i], dims[i+1], time_dim))
        self.down_samplers.append(DownSampler(dims[-2], dims[-1], time_dim, is_last = True))  # 3 -> 4

        # Middle Block:
        self.middle_block = MiddleBlock(dims[-1], time_dim)

        # Up Sampling:
        self.up_samplers = nn.ModuleList([])
        for i in range(len(dims) - 1, 1, -1):  # Loops 4, 3, 2  since the last one is a seperate convolution
            self.up_samplers.append(UpSampler(dims[i-1], dims[i], time_dim))

        # Final Convolution:
        self.final_conv = nn.Sequential(Conv1dBlock(dims[1], dims[1], kernel_size = 5),
                                        nn.Conv1d(dims[1], input_dim, kernel_size = 1))
        
        self.model_name = model_name
        if not os.path.exists(model_name):
            os.mkdir(model_name)
            self.losses = np.array([])
        else:
            self.load()

        _ = self.to(device)

    def forward(self, x, t):
        """
        x => Tensor of size (batch_size, traj_len*2)
        t => Integer representing the diffusion timestep of x
        """
        
        # Get the time embedding from t:
        time_emb = self.time_embedding(t)

        # Down Sampling Layers:
        h_list = []
        for i in range(len(self.down_samplers)):
            x, h = self.down_samplers[i](x, time_emb)
            h_list.append(h)

        # Middle Layer:
        x = self.middle_block(x, time_emb)

        # Up Sampling Layers:
        for i in range(len(self.up_samplers)):
            h_temp = h_list.pop()
            # print(f"Shape of x: {x.shape}\t Shape of h_list: {h_temp.shape}")
            x = self.up_samplers[i](x, h_temp, time_emb)   # How does pop work and not h_list[i]
            if x.shape[2] == 8 or x.shape[2] == 14 or x.shape[2] == 26 or x.shape[2] == 8: # Upsampling doubles the dimensions of the input. So, we are manually cropping the extra size to match the size of it's corresponding h (context/residual from the downsampling layer)
                x = tvtf.crop(x, 0, 0, x.shape[1], x.shape[2] - 1)

        # Final Convolution
        out = self.final_conv(x)

        return out

    def save(self):

        torch.save(self.state_dict(), self.model_name + "/weights_latest.pt")
        np.save(self.model_name + "/losses.npy", self.losses)

    def save_checkpoint(self, checkpoint):
        
        torch.save(self.state_dict(), self.model_name + "/weights_" + str(checkpoint) + ".pt")
        np.save(self.model_name + "/latest_checkpoint.npy", checkpoint)
    
    def load(self):

        self.losses = np.load(self.model_name + "/losses.npy")
        self.load_state_dict(torch.load(self.model_name + "/weights_latest.pt"))
        print("Loaded Model at " + str(self.losses.size) + " epochs")

    def load_checkpoint(self, checkpoint):

        _ = input("Press Enter if you are running the model for inference, or Ctrl+C\n(Never load a checkpoint for training! This will overwrite progress)")
        
        latest_checkpoint = np.load(self.model_name + "/latest_checkpoint.npy")
        self.load_state_dict(torch.load(self.model_name + "/weights_" + str(checkpoint) + ".pt"))
        self.losses = np.load(self.model_name + "/losses.npy")[:checkpoint]

    

            
diffusion/models/temporalunet.py
vishal-2000/EDMP
'''Guide Class

self.cfg is a dictionary. Example of self.cfg:

index: 1

hyperparameters:
  batch_size: 10
  obstacle_clearance:
    range: [0., 0.]
  obstacle_expansion:
    ind_set_range1: [255, 150]
    val_range1: [0.4, 0.4]
    ind_set_range2: [150, 20]
    val_range2: [0.4, 0.]
    ind_set_range3: [20, 0]
    val_range3: [0., 0.]
  guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume; 'ipv': interpolated_volume
  grad_norm: True
  guidance_period: 2
  guidance_stop_timestep: 5
  guidance_schedule:
    type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T))
    scale_val: 0.05
'''
from autolab_core import YamlConfig

class Guide:
    def __init__(self, path_to_yaml='./guides/cfgs/guide1.yaml') -> None:
        self.cfg = YamlConfig(path_to_yaml) # Dictionary

if __name__=='__main__':
    classifier = Guide()
    
guides/guide_cfg.py
vishal-2000/EDMP
from lib import *
from diffusion import *
from datasets.load_test_dataset import TestDataset

import argparse

from autolab_core import YamlConfig

import os
import wandb

import time

if __name__=='__main__':
    parser = argparse.ArgumentParser(
        prog='Benchmarking Diffusion',
        description='Benchmarking with IK on Test sets',
    )
    parser.add_argument('-c', '--cfg_path', type=str, default='./benchmark/cfgs/cfg1.yaml') # access as args.cfg_path


    args = parser.parse_args() 

    # Load config file
    benchmark_cfg = YamlConfig(args.cfg_path)

    # Load guide params
    guides = benchmark_cfg['guide']['guides'] # list
    guide_dpath = benchmark_cfg['guide']['guide_path'] 

    # Load model params
    device = benchmark_cfg['model']['device'] if torch.cuda.is_available() else 'cpu'
    traj_len = benchmark_cfg['model']['traj_len']
    T = benchmark_cfg['model']['T']
    num_channels = benchmark_cfg['model']['num_channels']

    # Load dataset
    dataset = TestDataset(benchmark_cfg['dataset']['dataset_type'], d_path=benchmark_cfg['dataset']['path']) # Dataset: global, hybrid, or both; d_path: path to dataset
    print(f"Benchmarking dataset: {benchmark_cfg['dataset']['dataset_type']} \tScene types: {benchmark_cfg['dataset']['scene_types']}\tguides: {benchmark_cfg['guide']['guides']}")

    # Load models
    # Load Models:
    env = RobotEnvironment(gui = benchmark_cfg['general']['gui'])
    diffuser = Diffusion(T=benchmark_cfg['model']['T'], device = device)
    model_name = benchmark_cfg['model']['model_dir'] + "TemporalUNetModel" + str(T) + "_N" + str(traj_len)
    if not os.path.exists(model_name):
        print("Model does not exist for these parameters. Train a model first.")
        _ = input("Press anything to exit")
        exit()
    denoiser = TemporalUNet(model_name = model_name, input_dim = num_channels, time_dim = 32, dims=(32, 64, 128, 256, 512, 512),
                            device = device)
    

    enable_wandb = benchmark_cfg['general']['wandb']['enable_wandb']
    

    guide_cfgs = []

    for g_index in guides:
        print(f"Loading Guide {g_index}" + "+-+-"*10)
        g_cfg = YamlConfig(benchmark_cfg['guide']['guide_path'] + f'cfgs/guide{g_index}.yaml')
        obstacle_clearance = np.linspace(g_cfg['hyperparameters']['obstacle_clearance']['range'][0], g_cfg['hyperparameters']['obstacle_clearance']['range'][1], 255)
        expansion_schedule = np.ones(shape = (255,))
        o_e_cfg = g_cfg['hyperparameters']['obstacle_expansion']
        expansion_schedule[o_e_cfg['isr1'][0]: o_e_cfg['isr1'][1]] = np.linspace(o_e_cfg['val1'][0], o_e_cfg['val1'][1], num=abs(o_e_cfg['isr1'][1] - o_e_cfg['isr1'][0]))
        expansion_schedule[o_e_cfg['isr2'][0]: o_e_cfg['isr2'][1]] = np.linspace(o_e_cfg['val2'][0], o_e_cfg['val2'][1], num=abs(o_e_cfg['isr2'][1] - o_e_cfg['isr2'][0]))
        expansion_schedule[o_e_cfg['isr3'][0]: o_e_cfg['isr3'][1]] = np.linspace(o_e_cfg['val3'][0], o_e_cfg['val3'][1], num=abs(o_e_cfg['isr3'][1] - o_e_cfg['isr3'][0]))
        volume_trust_region = g_cfg['hyperparameters']['volume_trust_region']

        guide_cfg = {
            'index': g_index,
            'clearance': obstacle_clearance, 
            'expansion': expansion_schedule, 
            'guidance_method': g_cfg['hyperparameters']['guidance_method'], 
            'grad_norm': g_cfg['hyperparameters']['grad_norm'],
            'volume_trust_region': g_cfg['hyperparameters']['volume_trust_region']
        }
        guide_cfgs.append(guide_cfg)
   

    t_success = 0
    for scene_type in benchmark_cfg['dataset']['scene_types']:
        print(f"Loading Scene Type: {scene_type}" + "+-+-"*10)
        i = 0
        for scene_num in range(dataset.data_nums[scene_type]):
            print(f"Scene num: {i+1}\tSuccess_rate: {t_success}/{i}")
            env.clear_obstacles()
            env.go_home()

            obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders, start_joints, all_ik_goals = dataset.fetch_data(scene_num=scene_num, scene_type=scene_type)

            guide_set = []
            final_trajectories = []

            start_time = time.time()
            st_time2 = time.time()
            for guide_cfg in guide_cfgs:
                print(f"Guide: {guide_cfg['index']}, Previous planning time: {time.time() - st_time2}") #, end = "")
                st_time2 = time.time()
                guide = IntersectionVolumeGuide(obstacle_config = obstacle_config, device = device, clearance = guide_cfg['clearance'], expansion = guide_cfg['expansion'], 
                                                guidance_method=guide_cfg['guidance_method'], grad_norm=guide_cfg['grad_norm'])
                metrics_calculator = MetricsCalculator(guide)
                guide_set.append([guide, metrics_calculator])

                # Filter final IKs
                st3 = time.time()
                volumes = guide.cost(torch.tensor(all_ik_goals.reshape((-1, 7, 1)), device=device), 0).sum(axis=(1, 2)).cpu().numpy()
                # Sort joints and volumes in the increasing order of volumes
                min_volume = np.min(volumes)
                indices = np.argsort(volumes)#[:10]
                rearranged_volumes = volumes[indices]
                goal_joints = all_ik_goals[indices]
                goal_joints = goal_joints[rearranged_volumes < min_volume + volume_trust_region]
                ideal_ind = np.argmin(np.linalg.norm(start_joints - goal_joints, axis=1))

                goal_joints = goal_joints[ideal_ind]
                print(f"IK time: {time.time() - st3}")

                st4 = time.time()

                trajectories = diffuser.denoise_guided(model = denoiser,
                                                        guide = guide,
                                                        batch_size = g_cfg['hyperparameters']['batch_size'],
                                                        traj_len = traj_len,
                                                        num_channels = num_channels,
                                                        condition = True,
                                                        benchmarking = True,
                                                        start = start_joints,
                                                        goal = goal_joints,
                                                        guidance_schedule=g_cfg['hyperparameters']['guidance_schedule']['type'],
                                                        guidance_scale=g_cfg['hyperparameters']['guidance_schedule']['scale_val'])
                                
                print(f"Denoiser time: {time.time() - st4}")

                trajectory = guide.choose_best_trajectory(start_joints, goal_joints, trajectories)
                final_trajectories.append(trajectory.copy())
                # trajectory is (7, 50) numpy array
                end_time = time.time()

            main_guide = IntersectionVolumeGuide(obstacle_config = obstacle_config, device = device, clearance = guide_cfgs[0]['clearance'], expansion = guide_cfgs[0]['expansion'], 
                                                guidance_method=guide_cfgs[0]['guidance_method'], grad_norm=guide_cfgs[0]['grad_norm'])
            
            trajectory = main_guide.choose_best_trajectory_final(np.array(final_trajectories))

            print(f"\nPlanning Time: {time.time() - start_time} seconds\n")

            if num_cuboids > 0:
                env.spawn_collision_cuboids(cuboid_config)
            if num_cylinders > 0:
                env.spawn_collision_cylinders(cylinder_config)

            
            success = env.benchmark_trajectory(trajectory)
            
            t_success += success
            # print(f"Success: {success}")

            i+=1


    
infer_serial.py
vishal-2000/EDMP
from .environment import *
from .guide import *
from .metrics import *
lib/__init__.py
vishal-2000/EDMP
import pybullet as p
import pybullet_utils.bullet_client as bc
import pybullet_data

# from mpinets.types import PlanningProblem, ProblemSet
# from geometrout.primitive import Cuboid, Cylinder
from robofin.robots import FrankaRobot

import numpy as np
import time
import re
import os
import h5py

class RobotEnvironment:

    def __init__(self, gui = True, timestep = 1/480, manipulator = True, benchmarking = False):

        self.client_id = bc.BulletClient(p.GUI if gui else p.DIRECT)            # Initialize the bullet client
        self.client_id.setAdditionalSearchPath(pybullet_data.getDataPath())     # Add pybullet's data package to path
        self.client_id.setTimeStep(timestep)                                    # Set simulation timestep
        self.client_id.configureDebugVisualizer(p.COV_ENABLE_SHADOWS, 0)        # Disable Shadows

        self.timestep = timestep
        self.benchmarking = benchmarking
        
        p.setAdditionalSearchPath(pybullet_data.getDataPath())      # Add pybullet's data package to path   

        self.colors = {'blue': np.array([78, 121, 167]) / 255.0,  # blue
                       'green': np.array([89, 161, 79]) / 255.0,  # green
                       'brown': np.array([156, 117, 95]) / 255.0,  # brown
                       'orange': np.array([242, 142, 43]) / 255.0,  # orange
                       'yellow': np.array([237, 201, 72]) / 255.0,  # yellow
                       'gray': np.array([186, 176, 172]) / 255.0,  # gray
                       'red': np.array([255, 87, 89]) / 255.0,  # red
                       'purple': np.array([176, 122, 161]) / 255.0,  # purple
                       'cyan': np.array([118, 183, 178]) / 255.0,  # cyan
                       'pink': np.array([255, 157, 167]) / 255.0}  # pink
        
        target = self.client_id.getDebugVisualizerCamera()[11]          # Get cartesian coordinates of the camera's focus
        self.client_id.resetDebugVisualizerCamera(                      # Reset initial camera position
            cameraDistance=1.5,
            cameraYaw=90,
            cameraPitch=-25,
            cameraTargetPosition=target,
        )
        
        p.resetSimulation()                             
        self.client_id.setGravity(0, 0, -9.8)           # Set Gravity

        # self.plane = self.client_id.loadURDF("plane.urdf", basePosition=(0, 0, 0), useFixedBase=True)   # Load a floor
        
        # self.client_id.changeDynamics(                  # Set physical properties of the floor
        #     self.plane,
        #     -1,
        #     lateralFriction=1.1,
        #     restitution=0.5,
        #     linearDamping=0.5,
        #     angularDamping=0.5,
        # )

        if manipulator:
            self.initialize_manipulator()

        self.obs_ids = []

    def initialize_manipulator(self, urdf_file = "franka_panda/panda.urdf", base_position = (0, 0, 0)):

        self.manipulator = self.client_id.loadURDF(urdf_file, basePosition = base_position, useFixedBase = True)
        self.joints = []

        for i in range(self.client_id.getNumJoints(self.manipulator)):

            info = self.client_id.getJointInfo(self.manipulator, i)

            joint_id = info[0]
            joint_name = info[1].decode("utf-8")
            joint_type = info[2]

            if joint_name == "panda_grasptarget_hand":
                self.end_effector = joint_id

            if joint_type == self.client_id.JOINT_REVOLUTE:
                self.joints.append(joint_id)

        self.joint_lower_limits = np.array([-166*(np.pi/180), 
                                           -101*(np.pi/180), 
                                           -166*(np.pi/180), 
                                           -176*(np.pi/180),
                                           -166*(np.pi/180),
                                           -1*(np.pi/180),
                                           -166*(np.pi/180)])
        
        self.joint_upper_limits = np.array([166*(np.pi/180), 
                                           101*(np.pi/180), 
                                           166*(np.pi/180), 
                                           -4*(np.pi/180),
                                           166*(np.pi/180),
                                           215*(np.pi/180),
                                           166*(np.pi/180)])
        
        links_folder_path = f'{pybullet_data.getDataPath()}/franka_panda/meshes/collision/'
        try:
            link_file_names = os.listdir(links_folder_path)
        except OSError as e:
            print(f"Error reading files in folder: {e}")
       
        self.link_meshes = {}
        self.link_dimensions = {}
        self.link_centers = {}

        self.link_index_to_name = ['link0', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'link7',
                                   'hand', 'finger', 'finger', 'finger', 'finger']

        for file_name in link_file_names:

            if file_name[-4:] == ".obj":
                vertices = []    
                link_name = file_name[:-4]
                with open(links_folder_path + file_name, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line.startswith('v '):
                            vertex = re.split(r'\s+', line)[1:4]
                            vertex = np.array([float(coord) for coord in vertex])
                            vertices.append(vertex)
                self.link_meshes[link_name] = np.array(vertices)
                max_point = np.max(self.link_meshes[link_name], axis = 0)
                min_point = np.min(self.link_meshes[link_name], axis = 0)
                self.link_dimensions[link_name] = max_point - min_point
                self.link_centers[link_name] = self.link_dimensions[link_name]/2 + min_point
    
    def get_mpinet_scene(self, index, data = 'train'):

        d_path = '/home/vishal/Volume_E/Active/Undergrad_research/CoRL2023/datasets/nvidia_MPI/mpinets_hybrid_training_data/val/'
        # with h5py.File("mpinet_dataset/" + data + ".hdf5", "r") as f:
        with h5py.File(d_path + data + ".hdf5", "r") as f:

            cuboid_centers = f['cuboid_centers'][index]
            cuboid_dims = f['cuboid_dims'][index]
            cuboid_quaternions = np.roll(f['cuboid_quaternions'][index], -1, axis = 1)

            cylinder_centers = f['cylinder_centers'][index]
            cylinder_heights = f['cylinder_heights'][index]
            cylinder_quaternions = np.roll(f['cylinder_quaternions'][index], -1, axis = 1)
            cylinder_radii = f['cylinder_radii'][index]

            num_cuboids = np.argmax(np.any(cuboid_dims == 0, axis = 1))
            num_cylinders = np.argmax(np.any(cylinder_heights == 0, axis = 1))

            cuboid_centers = cuboid_centers[:num_cuboids]
            cuboid_dims = cuboid_dims[:num_cuboids]
            cuboid_quaternions = cuboid_quaternions[:num_cuboids]

            cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1)
            
            cylinder_centers = cylinder_centers[:num_cylinders]
            cylinder_heights = cylinder_heights[:num_cylinders]
            cylinder_quaternions = cylinder_quaternions[:num_cylinders]
            cylinder_radii = cylinder_radii[:num_cylinders]

            cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1)

            cylinder_cuboid_dims = np.zeros((num_cylinders, 3))
            cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0]
            cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0]

            obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0)
            obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0)
            obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0)

            obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1)

            start = f['hybrid_solutions'][index][0, :]
            goal = f['hybrid_solutions'][index][-1, :]

            print(f"Actual Start: {start}")
            print(f"Actual Goal: {goal}")

            goal_ee = FrankaRobot.fk(goal, eff_frame="right_gripper")

            # goal_ee = data.target # SE3 Pose

            joint_7_samples = np.concatenate((np.random.uniform(-2.8973, 2.8973, 50), np.linspace(-2.8973, 2.8973, 50)))
            all_ik_goals = []
            for ik_sol_num, joint_ang in enumerate(joint_7_samples):
                ik_solutions = np.array(FrankaRobot.ik(goal_ee, joint_ang, 'right_gripper')).reshape((-1, 7))
                if len(ik_solutions)==0:
                    continue

                if len(all_ik_goals)==0:
                    all_ik_goals = ik_solutions.reshape((-1, 7))
                else:
                    all_ik_goals = np.vstack((all_ik_goals, ik_solutions[0].reshape((-1, 7)))) # n*7
                
            print(f"Total IK Solutions: {len(all_ik_goals)}")

            return obstacle_config, cuboid_config, cylinder_config, start, all_ik_goals # goal
    
    def spawn_cuboids(self, cuboid_config):

        for i in range(cuboid_config.shape[0]):
            
            vuid = self.client_id.createVisualShape(p.GEOM_BOX, 
                                        halfExtents = cuboid_config[i, 7:]/2,
                                        rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])]))

            obs_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 
                                                basePosition = cuboid_config[i, :3], 
                                                baseOrientation = cuboid_config[i, 3:7])
            
            self.obs_ids.append(obs_id)

    def spawn_cylinders(self, cylinder_config):

        for i in range(cylinder_config.shape[0]):
            
            vuid = self.client_id.createVisualShape(p.GEOM_CYLINDER, 
                                        radius = cylinder_config[i, 7],
                                        length = cylinder_config[i, 8],
                                        rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])]))

            obs_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 
                                                basePosition = cylinder_config[i, :3], 
                                                baseOrientation = cylinder_config[i, 3:7])
            
            self.obs_ids.append(obs_id)
    
    def spawn_collision_cuboids(self, cuboid_config):

        for i in range(cuboid_config.shape[0]):
            
            cuid = self.client_id.createCollisionShape(p.GEOM_BOX, 
                                                       halfExtents = cuboid_config[i, 7:]/2)
            
            vuid = self.client_id.createVisualShape(p.GEOM_BOX, 
                                                    halfExtents = cuboid_config[i, 7:]/2,
                                                    rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])]))

            obs_id = self.client_id.createMultiBody(baseMass = 0.,
                                                    baseCollisionShapeIndex = cuid,
                                                    baseVisualShapeIndex = vuid, 
                                                    basePosition = cuboid_config[i, :3], 
                                                    baseOrientation = cuboid_config[i, 3:7])
            
            self.obs_ids.append(obs_id)

    def spawn_collision_cylinders(self, cylinder_config):

        for i in range(cylinder_config.shape[0]):
            
            cuid = self.client_id.createCollisionShape(p.GEOM_CYLINDER, 
                                                       radius = cylinder_config[i, 7],
                                                       height = cylinder_config[i, 8])                                                   
            
            vuid = self.client_id.createVisualShape(p.GEOM_CYLINDER, 
                                                    radius = cylinder_config[i, 7],
                                                    length = cylinder_config[i, 8],
                                                    rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])]))

            obs_id = self.client_id.createMultiBody(baseMass = 0.,
                                                    baseCollisionShapeIndex = cuid, 
                                                    baseVisualShapeIndex = vuid, 
                                                    basePosition = cylinder_config[i, :3], 
                                                    baseOrientation = cylinder_config[i, 3:7])
            
            self.obs_ids.append(obs_id)
    
    def clear_obstacles(self):

        for id in self.obs_ids:
            self.client_id.removeBody(id)
        self.obs_ids = []
    
    def clip_joints(self, joints):

        return np.clip(joints, self.joint_lower_limits[np.newaxis, :, np.newaxis], self.joint_upper_limits[np.newaxis, :, np.newaxis])
    
    def draw_link_bounding_boxes(self):

        self.link_poses = []
        self.link_bounding_vertices = []
        
        self.link_bounding_objs = []

        for link_index in range(0, 11):     # The 12th link (i.e. 11th index) is the grasp target and is not needed   
                
            if link_index not in [8, 9]:    

                link_name = self.link_index_to_name[link_index+1]

                l, b, h = self.link_dimensions[link_name][0], self.link_dimensions[link_name][1], self.link_dimensions[link_name][2]

                if link_index == 7:
                    frame_pos, _ = self.client_id.getLinkState(self.manipulator, 7)[4:6]
                    _, frame_ori = self.client_id.getLinkState(self.manipulator, 10)[4:6]
                elif link_index != -1:
                    frame_pos, frame_ori = self.client_id.getLinkState(self.manipulator, link_index)[4:6]
                else:
                    frame_pos, frame_ori = self.client_id.getBasePositionAndOrientation(self.manipulator)
                world_transform = self.pose_to_transformation(np.array([*frame_pos, *frame_ori]))

                link_dimensions = self.link_dimensions[link_name].copy()
                if link_index == 10:
                    link_dimensions[1] *= 4

                world_link_center = (world_transform @ np.vstack((np.expand_dims(self.link_centers[link_name], 1), 1)))[:-1, 0]

            vertices = np.array([[-l/2, -b/2, -h/2],
                                [ l/2, -b/2, -h/2],
                                [ l/2,  b/2, -h/2],
                                [-l/2,  b/2, -h/2],
                                [-l/2, -b/2,  h/2],
                                [ l/2, -b/2,  h/2],
                                [ l/2,  b/2,  h/2],
                                [-l/2,  b/2,  h/2]])
            vertices = vertices + np.array([self.link_centers[link_name]])
            vertices = world_transform @ np.vstack((vertices.T, np.ones(8)))
            self.link_bounding_vertices.append(vertices.T[:, :-1])

            self.link_poses.append(np.array([*world_link_center, *frame_ori]))
            
            vuid = self.client_id.createVisualShape(p.GEOM_BOX, 
                                    halfExtents = link_dimensions/2,
                                    rgbaColor = np.hstack([self.colors['red'], np.array([1.0])]))
            
            obj_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 
                                                    basePosition = world_link_center, 
                                                    baseOrientation = frame_ori)
            
            self.link_bounding_objs.append(obj_id)
    
    def clear_bounding_boxes(self):

        for obj_id in self.link_bounding_objs:
            self.client_id.removeBody(obj_id)
    
    def get_joint_positions(self):

        return np.array([self.client_id.getJointState(self.manipulator, i)[0] for i in self.joints])
    
    def get_joint_velocities(self):

        return np.array([self.client_id.getJointState(self.manipulator, i)[1] for i in self.joints])

    def get_tf_mat(self, i, joint_angles):
        
        dh_params = np.array([[0, 0.333, 0, joint_angles[0]],
                    [0, 0, -np.pi / 2, joint_angles[1]],
                    [0, 0.316, np.pi / 2, joint_angles[2]],
                    [0.0825, 0, np.pi / 2, joint_angles[3]],
                    [-0.0825, 0.384, -np.pi / 2, joint_angles[4]],
                    [0, 0, np.pi / 2, joint_angles[5]],
                    [0.088, 0, np.pi / 2, joint_angles[6]],
                    [0, 0.107, 0, 0],
                    [0, 0, 0, -np.pi / 4],
                    [0.0, 0.1034, 0, 0]], dtype=np.float64)
        
        a = dh_params[i][0]
        d = dh_params[i][1]
        alpha = dh_params[i][2]
        theta = dh_params[i][3]
        q = theta

        return np.array([[np.cos(q), -np.sin(q), 0, a],
                        [np.sin(q) * np.cos(alpha), np.cos(q) * np.cos(alpha), -np.sin(alpha), -np.sin(alpha) * d],
                        [np.sin(q) * np.sin(alpha), np.cos(q) * np.sin(alpha), np.cos(alpha), np.cos(alpha) * d],
                        [0, 0, 0, 1]])

    def draw_frame(self, transform, scale_factor = 0.2):

        unit_axes_world = np.array([[scale_factor, 0, 0], 
                                    [0, scale_factor, 0], 
                                    [0, 0, scale_factor],
                                    [1, 1, 1]])
        axis_points = ((transform @ unit_axes_world)[:3, :]).T
        axis_center = transform[:3, 3]

        l1 = self.client_id.addUserDebugLine(axis_center, axis_points[0], self.colors['red'], lineWidth = 4)
        l2 = self.client_id.addUserDebugLine(axis_center, axis_points[1], self.colors['green'], lineWidth = 4)
        l3 = self.client_id.addUserDebugLine(axis_center, axis_points[2], self.colors['blue'], lineWidth = 4)

        frame_id = [l1, l2, l3]

        return frame_id[:]
    
    def remove_frame(self, frame_id):

        for id in frame_id:
            self.client_id.removeUserDebugItem(id)
        
    def inverse_of_transform(self, matrix):

        # Extract the rotation part and translation part of the matrix
        rotation_part = matrix[:3, :3]
        translation_part = matrix[:3, 3]
        
        # Calculate the inverse of the rotation part
        inverse_rotation = np.linalg.inv(rotation_part)
        
        # Calculate the new translation by applying the inverse rotation
        inverse_translation = -inverse_rotation.dot(translation_part)
        
        # Create the inverse transformation matrix
        inverse_matrix = np.zeros_like(matrix)
        inverse_matrix[:3, :3] = inverse_rotation
        inverse_matrix[:3, 3] = inverse_translation
        inverse_matrix[3, 3] = 1.0
        
        return inverse_matrix.copy()
    
    def forward_kinematics(self, joint_angles):

        T_EE = np.identity(4)
        for i in range(7 + 3):
            T_EE = T_EE @ self.get_tf_mat(i, joint_angles)

        return T_EE
    
    def get_jacobian(self, joint_angles):

        T_EE = self.get_ee_tf_mat(joint_angles)

        J = np.zeros((6, 10))
        T = np.identity(4)
        for i in range(7 + 3):
            T = T @ self.get_tf_mat(i, joint_angles)

            p = T_EE[:3, 3] - T[:3, 3]
            z = T[:3, 2]

            J[:3, i] = np.cross(z, p)
            J[3:, i] = z

        return J[:, :7]
    
    def inverse_kinematics(self, point, euler_orientation = None):

        if type(euler_orientation) == type(None):
            point = self.client_id.calculateInverseKinematics(self.manipulator, self.end_effector, point)
        else:
            quat = self.client_id.getQuaternionFromEuler(euler_orientation)
            point = self.client_id.calculateInverseKinematics(self.manipulator, self.end_effector, point, quat)

        return point
    
    def get_end_effector_pose(self):

        pos, ori = self.client_id.getLinkState(self.manipulator, 11, computeForwardKinematics = 1)[:2]
        pose = np.array([*pos, *ori])

        return pose
    
    def get_end_effector_transformation(self):

        pose = self.get_end_effector_pose()        
        transform = self.pose_to_transformation(pose)

        return transform

    def pose_to_transformation(self, pose):

        pos = pose[:3]
        quat = pose[3:]

        rotation_matrix = self.quaternion_to_rotation_matrix(quat)

        transform = np.zeros((4, 4))
        transform[:3, :3] = rotation_matrix.copy()
        transform[:3, 3] = pos.copy()
        transform[3, 3] = 1

        return transform

    def euler_to_rotation_matrix(yaw, pitch, roll):
        
        Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                    [np.sin(yaw), np.cos(yaw), 0],
                    [0, 0, 1]])

        Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                    [0, 1, 0],
                    [-np.sin(pitch), 0, np.cos(pitch)]])

        Rx = np.array([[1, 0, 0],
                    [0, np.cos(roll), -np.sin(roll)],
                    [0, np.sin(roll), np.cos(roll)]])

        R = Rz @ (Ry @ Rx)
        
        return R
    
    def quaternion_to_rotation_matrix(self, quat):
        """
        Convert a quaternion to a rotation matrix.
        
        :param q: Quaternion [w, x, y, z]
        :return: 3x3 rotation matrix
        """
        # w, x, y, z = quat
        # rotation_matrix = np.array([[1 - 2*y**2 - 2*z**2,  2*x*y - 2*z*w,        2*x*z + 2*y*w],
        #                             [2*x*y + 2*z*w,        1 - 2*x**2 - 2*z**2,  2*y*z - 2*x*w],
        #                             [2*x*z - 2*y*w,        2*y*z + 2*x*w,        1 - 2*x**2 - 2*y**2]])
        
        mat = np.array(self.client_id.getMatrixFromQuaternion(quat))
        rotation_matrix = np.reshape(mat, (3, 3))

        return rotation_matrix
        
    def rotation_matrix_to_quaternion(self, R):
        
        trace = np.trace(R)

        if trace > 0:
            S = np.sqrt(trace + 1.0) * 2
            qw = 0.25 * S
            qx = (R[2, 1] - R[1, 2]) / S
            qy = (R[0, 2] - R[2, 0]) / S
            qz = (R[1, 0] - R[0, 1]) / S
        elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
            S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2
            qw = (R[2, 1] - R[1, 2]) / S
            qx = 0.25 * S
            qy = (R[0, 1] + R[1, 0]) / S
            qz = (R[0, 2] + R[2, 0]) / S
        elif R[1, 1] > R[2, 2]:
            S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2
            qw = (R[0, 2] - R[2, 0]) / S
            qx = (R[0, 1] + R[1, 0]) / S
            qy = 0.25 * S
            qz = (R[1, 2] + R[2, 1]) / S
        else:
            S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2
            qw = (R[1, 0] - R[0, 1]) / S
            qx = (R[0, 2] + R[2, 0]) / S
            qy = (R[1, 2] + R[2, 1]) / S
            qz = 0.25 * S

        return np.array([qx, qy, qz, qw])
    
    def move_joints(self, target_joint_pos, speed = 0.01, timeout = 3):

        t0 = time.time()
        all_joints = []

        while (time.time() - t0) < timeout:
          
            current_joint_pos = self.get_joint_positions()
            error = target_joint_pos - current_joint_pos

            if all(np.abs(error) < 1e-2):
                for _ in range(10):
                    self.client_id.stepSimulation()     # Give time to stop
                    if self.benchmarking:
                        self.check_collisions()
                        if self.num_collisions > 0:
                            break
                if self.benchmarking and (self.num_collisions > 0):
                    break
                return True, all_joints
            
            norm = np.linalg.norm(error)
            vel = error / norm if norm > 0 else 0
            next_joint_pos = current_joint_pos + vel * speed
            all_joints.append(next_joint_pos)

            self.client_id.setJointMotorControlArray(           # Move with constant velocity
                bodyIndex = self.manipulator,
                jointIndices = self.joints,
                controlMode = p.POSITION_CONTROL,
                targetPositions = next_joint_pos,
                positionGains = np.ones(len(self.joints)),
            )

            self.client_id.stepSimulation()
            if self.benchmarking: 
                self.check_collisions()
                if self.num_collisions > 0:
                    break

        print(f"Warning: move_joints exceeded {timeout} second timeout. Skipping.")

        return False, all_joints
    
    def go_home(self):
        for i, joint_ind in enumerate(self.joints):
            self.client_id.resetJointState(self.manipulator, joint_ind, 0)
        # self.move_joints(np.zeros((7,)))

    def check_collisions(self):

        # There are 12 links.

        # Collision Info format:
        # ((0, 0, 7, 6, -1, (0.6231421349501027, -0.20546030368880996, 0.4579269918307523), 
        #   (0.6231421349501027, -0.20546030368881, 0.45900286827236414), (-4.511815760890017e-15, 2.2559078804450087e-14, -1.0), 
        #   0.0010758764416118338, 5678.289055003103, 185.98474464217657, (0.0, 1.0, 2.2559078804450087e-14), 1407.3361669376961, 
        #   (1.0, 1.0178240730107782e-28, -4.511815760890017e-15)),)

        self.num_timesteps += 1

        for obs_id in range(len(self.obs_ids)):

            info = self.client_id.getContactPoints(self.manipulator, self.obs_ids[obs_id]) #, link_index)
            if len(info) > 0:
                self.num_collisions += 1
                break
    
    def execute_trajectory(self, trajectory):

        for i, joint_ind in enumerate(self.joints):
            self.client_id.resetJointState(self.manipulator, joint_ind, trajectory[i, 0])

        _ = input("Press Enter to execute trajectory")

        for i in range(1, trajectory.shape[-1]):
            time.sleep(0.4)
            # current_joints = np.array([self.client_id.getJointState(self.manipulator, i)[0] for i in self.joints])
            target_joints = trajectory[:, i]
            # print(f"Current Joints: {current_joints}")
            # print(f"Target Joints: {target_joints}")
            # print(f"Itr number: {i}")

            if any(target_joints <= self.joint_lower_limits) or any(target_joints >= self.joint_upper_limits):

                print("Joint Limits Exceeded")

            self.move_joints(target_joints)
            self.client_id.stepSimulation()
    
    def benchmark_trajectory(self, trajectory): # , guide, metrics_calculator):

        self.benchmarking = True
        
        for i, joint_ind in enumerate(self.joints):
            self.client_id.resetJointState(self.manipulator, joint_ind, trajectory[i, 0])

        self.client_id.setJointMotorControlArray(           # Stabilize Robot
                bodyIndex = self.manipulator,
                jointIndices = self.joints,
                controlMode = p.POSITION_CONTROL,
                targetPositions = trajectory[:, 0],
                positionGains = np.zeros((len(self.joints))),
            )
        
        for i in range(10):
            self.client_id.stepSimulation()

        self.num_timesteps = 0
        self.num_collisions = 0

        dt = [0]

        for i in range(1, trajectory.shape[-1]):
            time.sleep(0.4)
            target_joints = trajectory[:, i]

            if any(target_joints <= self.joint_lower_limits) or any(target_joints >= self.joint_upper_limits):

                print("Joint Limits Exceeded")

            self.move_joints(target_joints)
            self.client_id.stepSimulation()
            self.check_collisions()
            if self.num_collisions > 0:
                break
            dt.append(self.num_timesteps * self.timestep)

        dt = np.diff(dt, n=1)[:, np.newaxis]

        collision_percentage = np.round((self.num_collisions / self.num_timesteps) * 100, 2)
        success = 1 if self.num_collisions == 0 else 0

        # joint_smoothness, end_eff_smoothness = metrics_calculator.smoothness_metric(trajectory, self.timestep)
        # joint_path_length, end_eff_path_length = metrics_calculator.path_length_metric(trajectory)

        self.benchmarking = False

        return success #, joint_path_length, end_eff_path_length, joint_smoothness[0], end_eff_smoothness[0]



        






lib/environment.py
vishal-2000/EDMP
import numpy as np
import torch
import os
import re
from einops.layers.torch import Rearrange
from scipy.spatial.transform import Rotation as R
# import pybullet as p

import pybullet_data

class IntersectionVolumeGuide:

    def __init__(self, obstacle_config, device, clearance, expansion, guidance_method='sv', grad_norm=True):
        '''
        Parameters:

        cost_type = 'iv' or 'sv' or 'ipv' (intersection volume or swept volume or interpolated volume) 
        grad_norm = True or False (indicates whether to take the norm of the gradient or not)
        '''

        # self.env = env
        self.device = device
        self.clearance = clearance
        self.expansion = expansion
        self.guidance_method = guidance_method
        self.grad_norm = grad_norm
        
        self.obstacle_config = np.array(obstacle_config)
        self.obs_ids = []

        self.static_dh_params = torch.tensor([[0, 0.333, 0, 0],
                                  [0, 0, -torch.pi / 2, 0],
                                  [0, 0.316, torch.pi / 2, 0],
                                  [0.0825, 0, torch.pi / 2, 0],
                                  [-0.0825, 0.384, -torch.pi / 2, 0],
                                  [0, 0, torch.pi / 2, 0],
                                  [0.088, 0, torch.pi / 2, 0],
                                  [0, 0.107, 0, 0],
                                  [0, 0, 0, -torch.pi / 4],
                                  [0.0, 0.1034, 0, 0]], dtype = torch.float32, device = self.device)
        
        self.define_link_information()
        # self.define_obstacles(obstacle_config)

        self.rearrange_joints = Rearrange('batch channels traj_len -> batch traj_len channels')
    
    def get_tf_mat(self, dh_params):
        
        # dh_params is (batch, traj_len, 1, 4)
        a = dh_params[:, :, 0]
        d = dh_params[:, :, 1]
        alpha = dh_params[:, :, 2]
        q = dh_params[:, :, 3]

        transform = torch.zeros(dh_params.shape[0], dh_params.shape[1], 4, 4)

        transform[:, :, 0, 0] = torch.cos(q)
        transform[:, :, 0, 1] = -torch.sin(q)
        transform[:, :, 0, 3] = a

        transform[:, :, 1, 0] = torch.sin(q) * torch.cos(alpha)
        transform[:, :, 1, 1] = torch.cos(q) * torch.cos(alpha)
        transform[:, :, 1, 2] = -torch.sin(alpha)
        transform[:, :, 1, 3] = -torch.sin(alpha) * d

        transform[:, :, 2, 0] = torch.sin(q) * torch.sin(alpha)
        transform[:, :, 2, 1] = torch.cos(q) * torch.sin(alpha)
        transform[:, :, 2, 2] = torch.cos(alpha)
        transform[:, :, 2, 3] = torch.cos(alpha) * d

        transform[:, :, -1, -1] = 1

        # Return transform of (batch, traj_len, 4, 4)
        return transform
    
    def forward_kinematics(self, joints):
        """
        joint_angles => Array of shape (batch, 7, traj_len)
        """        

        dh_params = torch.clone(self.static_dh_params)
        dh_params = dh_params.unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1)
        dh_params[:, :, :7, 3] = joints[:, :, :]
        # dh_params is (batch, traj_len, 10, 4)

        fk = torch.zeros(size = (dh_params.shape[0], dh_params.shape[1], 9, 4, 4), dtype = torch.float32, device = self.device)
        # fk is (batch, traj_len, 9, 4, 4)
    
        T = torch.eye(4).unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1)
        for i in range(7):
            dh_matrix = self.get_tf_mat(dh_params[:, :, i, :]) 
            # T is (batch, traj_len, 4, 4)
            # dh_matrix is (batch, traj_len, 4, 4)
            T = torch.matmul(T, dh_matrix) 
            if i == 6:
                fk[:, :, i:, :, :] = T.unsqueeze(2)
            else:
                fk[:, :, i, :, :] = T

        return fk
    
    def get_end_effector_transform(self, joints):

        # joints is (1, 50, 7) tensor

        dh_params = torch.clone(self.static_dh_params)
        dh_params = dh_params.unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1)
        dh_params[:, :, :7, 3] = joints[:, :, :]
        # dh_params is (1, 50, 10, 4)

        T = torch.eye(4).unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1)
        for i in range(10):
            dh_matrix = self.get_tf_mat(dh_params[:, :, i, :]) 
            # T is (batch, traj_len, 4, 4)
            # dh_matrix is (batch, traj_len, 4, 4)
            T = torch.matmul(T, dh_matrix)

        return T        

    def define_obstacles(self, obstacle_config, t):

        # Obstacle Config => (n, 10)
        
        obstacle_sizes = np.array(obstacle_config[:, 7:])
        obstacle_sizes = np.maximum(obstacle_sizes, self.expansion[t])

        if t != 0:
            obstacle_sizes = torch.tensor(obstacle_sizes, dtype = torch.float32, device = self.device) + self.clearance[t]
        else:
            obstacle_sizes = torch.tensor(obstacle_sizes, dtype = torch.float32, device = self.device)
        
        obstacle_static_vertices = self.get_vertices(obstacle_sizes)
        # obstacle_static_vertices => (n, 4, 8)

        obstacle_transform = np.zeros((obstacle_config.shape[0], 4, 4))
        for i in range(obstacle_config.shape[0]):
            obstacle_transform[i, :3, :3] = R.from_quat(obstacle_config[i, 3:7]).as_matrix() #).reshape((3, 3))
            obstacle_transform[i, :3, -1] = obstacle_config[i, :3]
        obstacle_transform[:, -1, -1] = 1.
        # obstacle_transform => (n, 4, 4)

        obstacle_transform = torch.tensor(obstacle_transform, dtype = torch.float32, device = self.device)
        obstacle_vertices = torch.bmm(obstacle_transform, obstacle_static_vertices)
        # obstacle_vertices => (n, 4, 8)

        self.obs_min = torch.min(obstacle_vertices, dim = -1)[0][:, :-1]
        self.obs_max = torch.max(obstacle_vertices, dim = -1)[0][:, :-1]
        # obs_min => (n, 3)
        # obs_max => (n, 3)
    
    def get_vertices(self, dimensions):

        l, b, h = dimensions[:, 0], dimensions[:, 1], dimensions[:, 2]
        obstacle_vertices = torch.zeros(size = (dimensions.shape[0], 4, 8), dtype = torch.float32, device = self.device)

        obstacle_vertices[:, 0, 0] = -l/2
        obstacle_vertices[:, 0, 1] = l/2
        obstacle_vertices[:, 0, 2] = l/2
        obstacle_vertices[:, 0, 3] = -l/2
        obstacle_vertices[:, 0, 4] = -l/2
        obstacle_vertices[:, 0, 5] = l/2
        obstacle_vertices[:, 0, 6] = l/2
        obstacle_vertices[:, 0, 7] = -l/2

        obstacle_vertices[:, 1, 0] = -b/2
        obstacle_vertices[:, 1, 1] = -b/2
        obstacle_vertices[:, 1, 2] = b/2
        obstacle_vertices[:, 1, 3] = b/2
        obstacle_vertices[:, 1, 4] = -b/2
        obstacle_vertices[:, 1, 5] = -b/2
        obstacle_vertices[:, 1, 6] = b/2
        obstacle_vertices[:, 1, 7] = b/2

        obstacle_vertices[:, 2, 0] = -h/2
        obstacle_vertices[:, 2, 1] = -h/2
        obstacle_vertices[:, 2, 2] = -h/2
        obstacle_vertices[:, 2, 3] = -h/2
        obstacle_vertices[:, 2, 4] = h/2
        obstacle_vertices[:, 2, 5] = h/2
        obstacle_vertices[:, 2, 6] = h/2
        obstacle_vertices[:, 2, 7] = h/2

        obstacle_vertices[:, 3, :] = 1.
        
        return obstacle_vertices
    
    def define_link_information(self):

        links_folder_path = f'{pybullet_data.getDataPath()}/franka_panda/meshes/collision/'
        try:
            link_file_names = os.listdir(links_folder_path)
        except OSError as e:
            print(f"Error reading files in folder: {e}")

        self.link_index_to_name = ['link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'link7',
                                   'hand', 'finger']
        self.link_dimensions_from_name = {}
        
        for file_name in link_file_names:

            if file_name[-4:] == ".obj":
                vertices = []    
                link_name = file_name[:-4]
                with open(links_folder_path + file_name, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line.startswith('v '):
                            vertex = re.split(r'\s+', line)[1:4]
                            vertex = np.array([float(coord) for coord in vertex])
                            vertices.append(vertex)
                max_point = np.max(np.array(vertices), axis = 0)
                min_point = np.min(np.array(vertices), axis = 0)
                self.link_dimensions_from_name[link_name] = max_point - min_point

        self.link_dimensions = []
        
        for link_index in range(len(self.link_index_to_name)):

            link_name = self.link_index_to_name[link_index]
            link_dimensions = self.link_dimensions_from_name[link_name].copy()
            
            if link_index == len(self.link_index_to_name) - 1:
                link_dimensions[1] *= 4
            
            self.link_dimensions.append(link_dimensions)
        self.link_dimensions = torch.tensor(np.array(self.link_dimensions), dtype = torch.float32, device = self.device)

        self.link_vertices = self.get_vertices(self.link_dimensions)

        self.link_static_joint_frame = [1, 2, 3, 4, 5, 6, 7, 7, 7] 
        self.static_frames = []

        # Link 0:
        self.static_frames.append([[1., 0., 0., 8.71e-05],
                                   [0., 1., 0., -3.709035e-02],
                                   [0., 0., 1., -6.851545e-02],
                                   [0., 0., 0., 1.]])
        # Link 1:
        self.static_frames.append([[1., 0., 0., -8.425e-05],
                                   [0., 1., 0., -6.93950016e-02],
                                   [0., 0., 1., 3.71961970e-02],
                                   [0., 0., 0., 1.]])
        
        # Link 2:
        self.static_frames.append([[1., 0., 0., 0.0414576],
                                   [0., 1., 0., 0.0281429],
                                   [0., 0., 1., -0.03293086],
                                   [0., 0., 0., 1.]])

        # Link 3:
        self.static_frames.append([[1., 0., 0., -4.12337575e-02],
                                   [0., 1., 0., 3.44296512e-02],
                                   [0., 0., 1., 2.79226985e-02],
                                   [0., 0., 0., 1.]])

        # Link 4:
        self.static_frames.append([[1., 0., 0., 3.3450000e-05],
                                   [0., 1., 0., 3.7388050e-02],
                                   [0., 0., 1., -1.0619285e-01],
                                   [0., 0., 0., 1.]])

        # Link 5:
        self.static_frames.append([[1., 0., 0., 4.21935000e-02],
                                   [0., 1., 0., 1.52195003e-02],
                                   [0., 0., 1., 6.07699933e-03],
                                   [0., 0., 0., 1.]])

        # Link 6:
        self.static_frames.append([[1., 0., 0., 1.86357500e-02],
                                   [0., 1., 0., 1.85788569e-02],
                                   [0., 0., 1., 7.94137484e-02],
                                   [0., 0., 0., 1.]])

        # Link 7:
        self.static_frames.append([[7.07106767e-01, 7.07106795e-01, 0., -1.26717073e-03],
                                   [-7.07106795e-01, 7.07106767e-01, 0., -1.25294673e-03],
                                   [0., 0., 1., 1.27018693e-01],
                                   [0., 0., 0., 1.]])

        # Link 8:
        self.static_frames.append([[7.07106767e-01, 7.07106795e-01, 0., 9.29352476e-03],
                                   [-7.07106795e-01, 7.07106767e-01, 0., 9.28272434e-03],
                                   [0., 0., 1., 1.92390375e-01],
                                   [0., 0., 0., 1.]])
        
        self.static_frames = torch.tensor(self.static_frames, dtype = torch.float32, device = self.device)
        
    def get_link_transform(self, joints):

        joint_transform = self.forward_kinematics(joints)

        # joint_transform => (batch, traj_len, 9, 4, 4)
        # static_frames => (9, 4, 4)
        link_transform = joint_transform @ self.static_frames.unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1, 1)

        return link_transform
    
    def cost(self, joint_input, t):

        # joint_angles => (b, 7, n)
        # obs_min => (n, 3)
        # obs_max => (n, 3) 
        self.define_obstacles(self.obstacle_config, t)

        joints = self.rearrange_joints(joint_input)
        # Now joints is (batch, traj_len, 7)
        
        link_transform = self.get_link_transform(joints)

        b = link_transform.shape[0]
        n = link_transform.shape[1]
        nl = link_transform.shape[2]
        no = self.obs_min.shape[0]

        # link_transform => (batch, traj_len, 9, 4, 4)
        # self.link_vertices => (9, 4, 8)
        link_vertices = link_transform @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(joints.shape[0], joints.shape[1], 1, 1, 1)
        link_vertices = link_vertices[:, :, :, :3, :]

        # link_vertices => (batch, traj_len, 9, 3, 8)
        link_min = torch.min(link_vertices, dim = -1)[0]
        link_max = torch.max(link_vertices, dim = -1)[0]

        # link_min => (batch, traj_len, 9, 3)
        # link_max => (batch, traj_len, 9, 3)
        expanded_link_min = link_min.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_max = link_max.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)

        expanded_obs_min = self.obs_min.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n, nl, 1, 1).view(b, n, no*nl, 3)
        expanded_obs_max = self.obs_max.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n, nl, 1, 1).view(b, n, no*nl, 3)

        overlap_min = torch.max(expanded_link_min, expanded_obs_min)
        overlap_max = torch.min(expanded_link_max, expanded_obs_max)

        overlap_lengths = overlap_max - overlap_min

        volumes = torch.prod(torch.clamp(overlap_lengths, min = 0), dim = -1)

        return volumes
    
    def interpolation_cost(self, joint_input, start, goal, t):

        # joint_angles => (b, 7, n)
        # obs_min => (n, 3)
        # obs_max => (n, 3) 
        self.define_obstacles(self.obstacle_config, t)

        joints = self.rearrange_joints(joint_input)
        # Now joints is (batch, traj_len, 7)
        
        joint_trajectory = torch.zeros(size = (joints.shape[0], joints.shape[1] + 2, joints.shape[2]))
        joint_trajectory[:, 0, :] = start.unsqueeze(0).repeat(joints.shape[0], 1)
        joint_trajectory[:, -1, :] = goal.unsqueeze(0).repeat(joints.shape[0], 1)
        joint_trajectory[:, 1:-1, :] = joints
        # Now joint_trajectory is (batch, traj_len + 2, 7)

        interpolated_joints1 = 0.5 * joint_trajectory[:, 1:, :] + 0.5 * joint_trajectory[:, :-1, :]
        interpolated_joints2 = 0.75 * joint_trajectory[:, 1:, :] + 0.25 * joint_trajectory[:, :-1, :]
        interpolated_joints3 = 0.25 * joint_trajectory[:, 1:, :] + 0.75 * joint_trajectory[:, :-1, :]
        
        link_transform1 = self.get_link_transform(interpolated_joints1)
        link_transform2 = self.get_link_transform(interpolated_joints2)
        link_transform3 = self.get_link_transform(interpolated_joints3)

        b = link_transform1.shape[0]
        n = link_transform1.shape[1]
        nl = link_transform1.shape[2]
        no = self.obs_min.shape[0]

        # link_transform => (batch, traj_len + 1, 9, 4, 4)
        # self.link_vertices => (9, 4, 8)
        link_vertices1 = link_transform1 @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(interpolated_joints1.shape[0], interpolated_joints1.shape[1], 1, 1, 1)
        link_vertices2 = link_transform2 @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(interpolated_joints1.shape[0], interpolated_joints1.shape[1], 1, 1, 1)
        link_vertices3 = link_transform3 @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(interpolated_joints1.shape[0], interpolated_joints1.shape[1], 1, 1, 1)
        
        link_vertices1 = link_vertices1[:, :, :, :3, :]
        link_vertices2 = link_vertices2[:, :, :, :3, :]
        link_vertices3 = link_vertices3[:, :, :, :3, :]

        # link_vertices => (batch, traj_len, 9, 3, 8)
        link_min1 = torch.min(link_vertices1, dim = -1)[0]
        link_max1 = torch.max(link_vertices1, dim = -1)[0]
        link_min2 = torch.min(link_vertices2, dim = -1)[0]
        link_max2 = torch.max(link_vertices2, dim = -1)[0]
        link_min3 = torch.min(link_vertices3, dim = -1)[0]
        link_max3 = torch.max(link_vertices3, dim = -1)[0]

        # link_min => (batch, traj_len, 9, 3)
        # link_max => (batch, traj_len, 9, 3)
        expanded_link_min1 = link_min1.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_max1 = link_max1.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_min2 = link_min2.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_max2 = link_max2.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_min3 = link_min3.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)
        expanded_link_max3 = link_max3.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n, no*nl, 3)

        expanded_obs_min = self.obs_min.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n, nl, 1, 1).view(b, n, no*nl, 3)
        expanded_obs_max = self.obs_max.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n, nl, 1, 1).view(b, n, no*nl, 3)

        overlap_min1 = torch.max(expanded_link_min1, expanded_obs_min)
        overlap_max1 = torch.min(expanded_link_max1, expanded_obs_max)
        overlap_min2 = torch.max(expanded_link_min2, expanded_obs_min)
        overlap_max2 = torch.min(expanded_link_max2, expanded_obs_max)
        overlap_min3 = torch.max(expanded_link_min3, expanded_obs_min)
        overlap_max3 = torch.min(expanded_link_max3, expanded_obs_max)

        overlap_lengths1 = overlap_max1 - overlap_min1
        overlap_lengths2 = overlap_max2 - overlap_min2
        overlap_lengths3 = overlap_max3 - overlap_min3

        volumes1 = torch.prod(torch.clamp(overlap_lengths1, min = 0), dim = -1)
        volumes2 = torch.prod(torch.clamp(overlap_lengths2, min = 0), dim = -1)
        volumes3 = torch.prod(torch.clamp(overlap_lengths3, min = 0), dim = -1)

        return (volumes1 + volumes2 + volumes3)
    
    def swept_volume_cost(self, joint_input, start, goal, t):

        # joint_angles => (b, 7, n)
        # obs_min => (n, 3)
        # obs_max => (n, 3) 
        self.define_obstacles(self.obstacle_config, t)

        joints = self.rearrange_joints(joint_input)
        # Now joints is (batch, traj_len, 7)

        joint_trajectory = torch.zeros(size = (joints.shape[0], joints.shape[1] + 2, joints.shape[2]))
        joint_trajectory[:, 0, :] = start.unsqueeze(0).repeat(joints.shape[0], 1)
        joint_trajectory[:, -1, :] = goal.unsqueeze(0).repeat(joints.shape[0], 1)
        joint_trajectory[:, 1:-1, :] = joints
        # Now joint_trajectory is (batch, traj_len + 2, 7)
        
        link_transform = self.get_link_transform(joint_trajectory)

        b = link_transform.shape[0]
        n = link_transform.shape[1]
        nl = link_transform.shape[2]
        no = self.obs_min.shape[0]

        # link_transform => (batch, traj_len+2, 9, 4, 4)
        # self.link_vertices => (9, 4, 8)
        link_vertices = link_transform @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(joint_trajectory.shape[0], joint_trajectory.shape[1], 1, 1, 1)
        link_vertices = link_vertices[:, :, :, :3, :]

        # link_vertices => (batch, traj_len+2, 9, 3, 8)
        link_min = torch.min(link_vertices, dim = -1)[0]
        link_max = torch.max(link_vertices, dim = -1)[0]

        # link_min => (batch, traj_len+2, 9, 3)
        # link_max => (batch, traj_len+2, 9, 3)
        link_min_A = link_min[:, :-1, :, :]
        link_max_A = link_max[:, :-1, :, :]
        link_min_B = link_min[:, 1:, :, :]
        link_max_B = link_max[:, 1:, :, :]
        swept_volumes_min = torch.min(link_min_A, link_min_B)
        swept_volumes_max = torch.max(link_max_A, link_max_B)

        # swept_volumes_min => (batch, traj_len+1, 9, 3)
        # swept_volumes_max => (batch, traj_len+1, 9, 3)
        expanded_link_min = swept_volumes_min.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n-1, no*nl, 3)
        expanded_link_max = swept_volumes_max.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n-1, no*nl, 3)

        expanded_obs_min = self.obs_min.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n-1, nl, 1, 1).view(b, n-1, no*nl, 3)
        expanded_obs_max = self.obs_max.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n-1, nl, 1, 1).view(b, n-1, no*nl, 3)

        overlap_min = torch.max(expanded_link_min, expanded_obs_min)
        overlap_max = torch.min(expanded_link_max, expanded_obs_max)

        overlap_lengths = overlap_max - overlap_min

        volumes = torch.prod(torch.clamp(overlap_lengths, min = 0), dim = -1)

        return volumes
    
    def swept_volume_cost_final(self, joint_input, t):

        # joint_angles => (b, 7, n)
        # obs_min => (n, 3)
        # obs_max => (n, 3) 
        self.define_obstacles(self.obstacle_config, t)

        joints = self.rearrange_joints(joint_input)
        # Now joints is (batch, traj_len, 7)

        joint_trajectory = torch.zeros(size = (joints.shape[0], joints.shape[1] + 2, joints.shape[2]))
        # joint_trajectory[:, 0, :] = start.unsqueeze(0).repeat(joints.shape[0], 1)
        # joint_trajectory[:, -1, :] = goal.unsqueeze(0).repeat(joints.shape[0], 1)
        joint_trajectory[:, 1:-1, :] = joints
        # Now joint_trajectory is (batch, traj_len + 2, 7)
        
        link_transform = self.get_link_transform(joint_trajectory)

        b = link_transform.shape[0]
        n = link_transform.shape[1]
        nl = link_transform.shape[2]
        no = self.obs_min.shape[0]

        # link_transform => (batch, traj_len+2, 9, 4, 4)
        # self.link_vertices => (9, 4, 8)
        link_vertices = link_transform @ self.link_vertices.unsqueeze(0).unsqueeze(0).repeat(joint_trajectory.shape[0], joint_trajectory.shape[1], 1, 1, 1)
        link_vertices = link_vertices[:, :, :, :3, :]

        # link_vertices => (batch, traj_len+2, 9, 3, 8)
        link_min = torch.min(link_vertices, dim = -1)[0]
        link_max = torch.max(link_vertices, dim = -1)[0]

        # link_min => (batch, traj_len+2, 9, 3)
        # link_max => (batch, traj_len+2, 9, 3)
        link_min_A = link_min[:, :-1, :, :]
        link_max_A = link_max[:, :-1, :, :]
        link_min_B = link_min[:, 1:, :, :]
        link_max_B = link_max[:, 1:, :, :]
        swept_volumes_min = torch.min(link_min_A, link_min_B)
        swept_volumes_max = torch.max(link_max_A, link_max_B)

        # swept_volumes_min => (batch, traj_len+1, 9, 3)
        # swept_volumes_max => (batch, traj_len+1, 9, 3)
        expanded_link_min = swept_volumes_min.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n-1, no*nl, 3)
        expanded_link_max = swept_volumes_max.unsqueeze(-2).repeat(1, 1, 1, no, 1).view(b, n-1, no*nl, 3)

        expanded_obs_min = self.obs_min.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n-1, nl, 1, 1).view(b, n-1, no*nl, 3)
        expanded_obs_max = self.obs_max.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(b, n-1, nl, 1, 1).view(b, n-1, no*nl, 3)

        overlap_min = torch.max(expanded_link_min, expanded_obs_min)
        overlap_max = torch.min(expanded_link_max, expanded_obs_max)

        overlap_lengths = overlap_max - overlap_min

        volumes = torch.prod(torch.clamp(overlap_lengths, min = 0), dim = -1)

        return volumes

    def get_gradient(self, joint_input, start, goal, t):

        joint_tensor1 = torch.tensor(joint_input, dtype = torch.float32, device = self.device)
        joint_tensor1.requires_grad = True
        joint_tensor1.grad = None

        start = torch.tensor(start, dtype = torch.float32, device = self.device)
        goal = torch.tensor(goal, dtype = torch.float32, device = self.device)
        # start => (7,)
        # goal => (7,)

        # DIFFERENT COSTS:
        # cost = torch.sum(self.cost(joint_tensor1, t)) # Intersection Volume
        # cost = torch.sum(self.swept_volume_cost(joint_tensor1, start, goal, t)) # Swept Volume
        # cost = torch.sum(self.cost(joint_tensor1, t)) + torch.sum(self.interpolation_cost(joint_tensor1, start, goal, t)) # Interpolation Volume

        if self.guidance_method =='iv':
            cost = torch.sum(self.cost(joint_tensor1, t)) # Intersection Volume
        elif self.guidance_method == 'sv':
            cost = torch.sum(self.swept_volume_cost(joint_tensor1, start, goal, t)) # Swept Volume
        elif self.guidance_method == 'ipv':
            cost = torch.sum(self.cost(joint_tensor1, t)) + torch.sum(self.interpolation_cost(joint_tensor1, start, goal, t)) # Interpolation Volume
        
        cost.backward()

        gradient1 = joint_tensor1.grad.cpu().numpy()

        if self.grad_norm == True:
            if np.linalg.norm(gradient1) > 0:    
                gradient1 = gradient1 / np.linalg.norm(gradient1)

        return gradient1
    
    def choose_best_trajectory(self, start, goal, trajectories):

        joint_tensor = torch.tensor(trajectories[:, :, 1:-1], dtype = torch.float32, device = self.device)
        
        start = torch.tensor(start, dtype = torch.float32, device = self.device)
        goal = torch.tensor(goal, dtype = torch.float32, device = self.device)
        
        overlap_volumes = torch.sum(self.swept_volume_cost(joint_tensor, start, goal, 0), dim = (1, 2))
        min_index = torch.argmin(overlap_volumes)

        # print(overlap_volumes)
        # print("best = ", overlap_volumes[min_index])

        return trajectories[min_index] #, min_index
    
    def choose_best_trajectory_final(self, trajectories):

        joint_tensor = torch.tensor(trajectories[:, :, 1:-1], dtype = torch.float32, device = self.device)
        
        # start = torch.tensor(start, dtype = torch.float32, device = self.device)
        # goal = torch.tensor(goal, dtype = torch.float32, device = self.device)
        
        overlap_volumes = torch.sum(self.swept_volume_cost_final(joint_tensor, 0), dim = (1, 2))
        min_index = torch.argmin(overlap_volumes)

        # print(overlap_volumes)
        # print("best = ", overlap_volumes[min_index])

        return trajectories[min_index] #, min_index

    def smoothness_cost(self, joints, start, goal):

        start = torch.tensor(start, dtype = torch.float32, device = self.device).view(1, 7, 1)
        goal = torch.tensor(goal, dtype = torch.float32, device = self.device).view(1, 7, 1)

        cost = torch.sum((joints[:, :, 0] - start)**2) + torch.sum((joints[:, :, 2:-1] - joints[:, :, 1:-2])**2) + torch.sum((goal - joints[:, :, -1])**2)

        return cost





lib/guide.py
vishal-2000/EDMP
import torch    
import numpy as np

class MetricsCalculator:
       
    def __init__(self, guide):

        self.guide = guide
        self.device = self.guide.device

    def smoothness_metric(self, joints, dt):
        """
        dt is time taken between each step
        joints is (7, 50) numpy array
        """

        joint_tensor = self.guide.rearrange_joints(torch.tensor(joints, dtype = torch.float32, device = self.device).unsqueeze(0))
        end_eff_transforms = self.guide.get_end_effector_transform(joint_tensor)

        end_eff_positions = (end_eff_transforms[0, :, :3, 3]).numpy(force=True)
        # end_eff_positions is (50, 3) numpy array

        reshaped_joints = joints.T
        # reshaped_joints is (50, 7) numpy array
        joint_smoothness = np.linalg.norm(np.diff(reshaped_joints, n=1, axis=0) / dt, axis=1)
        joint_smoothness = self.sparc(joint_smoothness, 1. / dt)
        
        end_eff_smoothness = np.linalg.norm(np.diff(end_eff_positions, n=1, axis=0) / dt, axis=1)
        end_eff_smoothness = self.sparc(end_eff_smoothness, 1. / dt)

        return joint_smoothness, end_eff_smoothness
    
    def path_length_metric(self, joints):

        joint_tensor = self.guide.rearrange_joints(torch.tensor(joints, dtype = torch.float32, device = self.device).unsqueeze(0))
        end_eff_transforms = self.guide.get_end_effector_transform(joint_tensor)
        end_eff_positions = (end_eff_transforms[0, :, :3, 3]).numpy(force=True)
        # end_eff_positions is (50, 3) numpy array

        reshaped_joints = joints.T

        end_eff_path_length = np.sum(np.linalg.norm(np.diff(end_eff_positions, 1, axis=0), axis=1))
        joint_path_length = np.sum(np.linalg.norm(np.diff(reshaped_joints, 1, axis=0), axis=1))

        return joint_path_length, end_eff_path_length

    def sparc(self, movement, fs, padlevel=4, fc=10.0, amp_th=0.05):
        """
        Calculates the smoothness of the given speed profile using the modified
        spectral arc length metric.
        Parameters
        ----------
        movement : np.array
                The array containing the movement speed profile.
        fs       : float
                The sampling frequency of the data.
        padlevel : integer, optional
                Indicates the amount of zero padding to be done to the movement
                data for estimating the spectral arc length. [default = 4]
        fc       : float, optional
                The max. cut off frequency for calculating the spectral arc
                length metric. [default = 10.]
        amp_th   : float, optional
                The amplitude threshold to used for determing the cut off
                frequency upto which the spectral arc length is to be estimated.
                [default = 0.05]
        Returns
        -------
        sal      : float
                The spectral arc length estimate of the given movement's
                smoothness.
        (f, Mf)  : tuple of two np.arrays
                This is the frequency(f) and the magntiude spectrum(Mf) of the
                given movement data. This spectral is from 0. to fs/2.
        (f_sel, Mf_sel) : tuple of two np.arrays
                        This is the portion of the spectrum that is selected for
                        calculating the spectral arc length.
        Notes
        -----
        This is the modfieid spectral arc length metric, which has been tested only
        for discrete movements.

        Examples
        --------
        >>> t = np.arange(-1, 1, 0.01)
        >>> move = np.exp(-5*pow(t, 2))
        >>> sal, _, _ = sparc(move, fs=100.)
        >>> '%.5f' % sal
        '-1.41403'
        """
        if np.allclose(movement, 0):
            print("All movement was 0, returning 0")
            return 0, None, None
        # Number of zeros to be padded.
        nfft = int(pow(2, np.ceil(np.log2(len(movement))) + padlevel))

        # Frequency
        f = np.arange(0, fs, fs / nfft)
        # Normalized magnitude spectrum
        Mf = abs(np.fft.fft(movement, nfft))
        Mf = Mf / max(Mf)

        # Indices to choose only the spectrum within the given cut off frequency
        # Fc.
        # NOTE: This is a low pass filtering operation to get rid of high frequency
        # noise from affecting the next step (amplitude threshold based cut off for
        # arc length calculation).
        fc_inx = ((f <= fc) * 1).nonzero()
        f_sel = f[fc_inx]
        Mf_sel = Mf[fc_inx]

        # Choose the amplitude threshold based cut off frequency.
        # Index of the last point on the magnitude spectrum that is greater than
        # or equal to the amplitude threshold.
        inx = ((Mf_sel >= amp_th) * 1).nonzero()[0]
        fc_inx = range(inx[0], inx[-1] + 1)
        f_sel = f_sel[fc_inx]
        Mf_sel = Mf_sel[fc_inx]

        # Calculate arc length
        new_sal = -sum(
            np.sqrt(
                pow(np.diff(f_sel) / (f_sel[-1] - f_sel[0]), 2) + pow(np.diff(Mf_sel), 2)
            )
        )
        return new_sal, (f, Mf), (f_sel, Mf_sel)
lib/metrics.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

from typing import Union, Sequence
import numpy as np
import random
import torch
from geometrout.primitive import Sphere, Cuboid, Cylinder


class TorchSpheres:
    """
    A Pytorch representation of a batch of M spheres (i.e. B elements in the batch,
    M spheres per element). Any of these spheres can have zero volume (these
    will be masked out during calculation of the various functions in this
    class, such as sdf).
    """

    def __init__(self, centers: torch.Tensor, radii: torch.Tensor):
        """
        :param centers torch.Tensor: a set of centers, has dim [B, M, 3]
        :param radii torch.Tensor: a set of radii, has dim [B, M, 1]
        """
        assert centers.ndim == 3
        assert radii.ndim == 3

        # TODO It would be more memory efficient to rely more heavily on broadcasting
        # in some cases where multiple spheres have the same size
        assert centers.ndim == radii.ndim

        # This logic is to determine the batch size. Either batch sizes need to
        # match or if only one variable has a batch size, that one is assumed to
        # be the batch size

        self.centers = centers
        self.radii = radii
        self.mask = ~torch.isclose(self.radii, torch.zeros(1).type_as(centers)).squeeze(
            -1
        )

    def surface_area(self) -> torch.Tensor:
        """
        Calculates the surface area of the spheres

        :rtype torch.Tensor: A tensor of the surface areas of the spheres
        """
        area = 4 * np.pi * torch.pow(self.radii, 3)
        return area

    def sample_surface(self, num_points: int) -> torch.Tensor:
        """
        Samples points from all spheres, including ones with zero volume

        :param num_points int: The number of points to sample per sphere
        :rtype torch.Tensor: The points, has dim [B, M, N]
        """
        B, M, _ = self.centers.shape
        unnormalized_points = torch.rand((B, M, num_points, 3))
        normalized = (
            unnormalized_points
            / torch.linalg.norm(unnormalized_points, dim=-1)[:, :, :, None]
        )
        random_points = (
            normalized * self.radii[:, :, None, :] + self.centers[:, :, None, :]
        )
        return random_points

    def sdf(self, points: torch.Tensor) -> torch.Tensor:
        """
        :param points torch.Tensor: The points with which to calculate the
                                    SDF, has dim [B, N, 3] (N is the number of points)
        :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF
                             value for each of the M spheres), has dim [B, N]
        """
        assert points.ndim == 3
        B, M, _ = self.radii.shape
        _, N, _ = points.shape
        all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points)
        distances = points[:, None, :, :] - self.centers[:, :, None, :]
        all_sdfs[self.mask] = (
            torch.linalg.norm(distances[self.mask], dim=-1) - self.radii[self.mask]
        )
        return torch.min(all_sdfs, dim=1)[0]

    def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor:
        """
        Calculates SDF values for a time sequence of point clouds
        :param points torch.Tensor: The batched sequence of point clouds with
                                    dimension [B, T, N, 3] (T in sequence length,
                                    N is number of points)
        :rtype torch.Tensor: The scene SDF for each point at each timestep (i.e. the minimum
                             SDF value across the M spheres at each timestep),
                             has dim [B, T, N]
        """
        assert points.ndim == 4
        B, M, _ = self.radii.shape
        _, T, N, _ = points.shape
        all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points)
        distances = points[:, None, :, :] - self.centers[:, :, None, None, :]
        all_sdfs[self.mask] = (
            torch.linalg.norm(distances[self.mask], dim=-1)
            - self.radii[:, :, None, :][self.mask]
        )
        return torch.min(all_sdfs, dim=1)[0]


class TorchCuboids:
    """
    A Pytorch representation of a batch of M cuboids (i.e. B elements in the batch,
    M cuboids per element). Any of these cuboids can have zero volume (these
    will be masked out during calculation of the various functions in this
    class, such as sdf).
    """

    def __init__(
        self, centers: torch.Tensor, dims: torch.Tensor, quaternions: torch.Tensor
    ):
        """
        :param centers torch.Tensor: Has dim [B, M, 3]
        :param dims torch.Tensor: Has dim [B, M, 3]
        :param quaternions torch.Tensor: Has dim [B, M, 4] with quaternions formatted as
                                   w, x, y, z
        """
        assert centers.ndim == 3
        assert dims.ndim == 3
        assert quaternions.ndim == 3

        self.dims = dims
        self.centers = centers

        # It's helpful to ensure the quaternions are normalized
        self.quats = quaternions / torch.linalg.norm(quaternions, dim=2)[:, :, None]

        self._init_frames()
        # Mask for nonzero volumes
        self.mask = ~torch.any(
            torch.isclose(self.dims, torch.zeros(1).type_as(centers)), dim=-1
        )

    def geometrout(self):
        """
        Helper method to convert this into geometrout primitives
        """
        B, M, _ = self.centers.shape
        return [
            [
                Cuboid(
                    center=self.centers[bidx, midx, :].detach().cpu().numpy(),
                    dims=self.dims[bidx, midx, :].detach().cpu().numpy(),
                    quaternion=self.quats[bidx, midx, :].detach().cpu().numpy(),
                )
                for midx in range(M)
                if self.mask[bidx, midx]
            ]
            for bidx in range(B)
        ]

    def _init_frames(self):
        """
        In order to calculate the SDF, we need to calculate the inverse
        transformation of the cuboid. This is because we are transforming points
        in the world frame into the cuboid frame.
        """

        # Initialize the inverse rotation
        w = self.quats[:, :, 0]
        x = -self.quats[:, :, 1]
        y = -self.quats[:, :, 2]
        z = -self.quats[:, :, 3]

        # Naming is a little disingenuous here because I'm multiplying everything by two,
        # but can't put numbers at the beginning of variable names.
        xx = 2 * torch.pow(x, 2)
        yy = 2 * torch.pow(y, 2)
        zz = 2 * torch.pow(z, 2)

        wx = 2 * w * x
        wy = 2 * w * y
        wz = 2 * w * z
        xy = 2 * x * y
        xz = 2 * x * z
        yz = 2 * y * z

        B, M, _ = self.centers.shape
        B = self.centers.size(0)
        M = self.centers.size(1)
        self.inv_frames = torch.zeros((B, M, 4, 4)).type_as(self.centers)
        self.inv_frames[:, :, 3, 3] = 1

        R = torch.stack(
            [
                torch.stack([1 - yy - zz, xy - wz, xz + wy], dim=2),
                torch.stack([xy + wz, 1 - xx - zz, yz - wx], dim=2),
                torch.stack([xz - wy, yz - wx, 1 - xx - yy], dim=2),
            ],
            dim=2,
        )
        Rt = torch.matmul(R, -1 * self.centers.unsqueeze(3)).squeeze(3)

        # Fill in the rotation matrices
        self.inv_frames[:, :, :3, :3] = R

        # Invert the transform by multiplying the inverse translation by the inverse rotation
        self.inv_frames[:, :, :3, 3] = Rt

    def surface_area(self) -> torch.Tensor:
        """
        Calculates the surface area of the cuboids

        :rtype torch.Tensor: A tensor of the surface areas of the cuboids
        """
        area = 2 * (
            self.dims[:, :, 0] * self.dims[:, :, 1]
            + self.dims[:, :, 0] * self.dims[:, :, 2]
            + self.dims[:, :, 1] * self.dims[:, :, 2]
        )
        return area

    def sdf(self, points: torch.Tensor) -> torch.Tensor:
        """
        :param points torch.Tensor: The points with which to calculate the SDF, has
                                    dim [B, N, 3] (N is the number of points)
        :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF
                             value for each of the M cuboids), has dim [B, N]
        """
        assert points.ndim == 3
        # We are going to need to map points in the global frame to the cuboid frame
        # First take the points and make them homogeneous by adding a one to the end

        # points_from_volumes = points[self.nonzero_volumes, :, :]
        assert points.size(0) == self.centers.size(0)
        if torch.all(~self.mask):
            return float("inf") * torch.ones(points.size(0), points.size(1)).type_as(
                points
            )

        homog_points = torch.cat(
            (
                points,
                torch.ones((points.size(0), points.size(1), 1)).type_as(points),
            ),
            dim=2,
        )
        # Next, project these points into their respective cuboid frames
        # Will return [B x M x N x 3]

        points_proj = torch.matmul(
            self.inv_frames[:, :, None, :, :], homog_points[:, None, :, :, None]
        ).squeeze(-1)[:, :, :, :3]
        B, M, N, _ = points_proj.shape
        masked_points = points_proj[self.mask]

        # The follow computations are adapted from here
        # https://github.com/fogleman/sdf/blob/main/sdf/d3.py
        # Move points to top corner

        distances = torch.abs(masked_points) - (self.dims[self.mask] / 2)[:, None, :]
        # This is distance only for points outside the box, all points inside return zero
        # This probably needs to be fixed or else there will be a nan gradient

        outside = torch.linalg.norm(
            torch.maximum(distances, torch.zeros_like(distances)), dim=-1
        )
        # This is distance for points inside the box, all others return zero
        inner_max_distance = torch.max(distances, dim=-1).values
        inside = torch.minimum(inner_max_distance, torch.zeros_like(inner_max_distance))
        all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points)
        all_sdfs[self.mask] = outside + inside
        return torch.min(all_sdfs, dim=1)[0]

    def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor:
        """
        Calculates SDF values for a time sequence of point clouds
        :param points torch.Tensor: The batched sequence of point clouds with
                                    dimension [B, T, N, 3] (T in sequence length,
                                    N is number of points)
        :rtype torch.Tensor: The scene SDF for each point at each timestep
                             (i.e. the minimum SDF value across the M cuboids
                             at each timestep), has dim [B, T, N]
        """
        assert points.ndim == 4

        # We are going to need to map points in the global frame to the cuboid frame
        # First take the points and make them homogeneous by adding a one to the end

        # points_from_volumes = points[self.nonzero_volumes, :, :]
        assert points.size(0) == self.centers.size(0)
        if torch.all(~self.mask):
            return float("inf") * torch.ones(points.shape[:-1]).type_as(points)

        homog_points = torch.cat(
            (
                points,
                torch.ones((*points.shape[:-1], 1)).type_as(points),
            ),
            dim=3,
        )
        # Next, project these points into their respective cuboid frames
        # Will return [B x M x N x 3]

        points_proj = torch.matmul(
            self.inv_frames[:, :, None, None, :, :],
            homog_points[:, None, :, :, :, None],
        ).squeeze(-1)[:, :, :, :, :3]
        B, M, T, N, _ = points_proj.shape
        assert T == points.size(1)
        assert N == points.size(2)
        masked_points = points_proj[self.mask]

        # The follow computations are adapted from here
        # https://github.com/fogleman/sdf/blob/main/sdf/d3.py
        # Move points to top corner

        distances = (
            torch.abs(masked_points) - (self.dims[self.mask] / 2)[:, None, None, :]
        )
        # This is distance only for points outside the box, all points inside return zero
        # This probably needs to be fixed or else there will be a nan gradient

        outside = torch.linalg.norm(
            torch.maximum(distances, torch.zeros_like(distances)), dim=-1
        )
        # This is distance for points inside the box, all others return zero
        inner_max_distance = torch.max(distances, dim=-1).values
        inside = torch.minimum(inner_max_distance, torch.zeros_like(inner_max_distance))
        all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points)
        all_sdfs[self.mask] = outside + inside
        return torch.min(all_sdfs, dim=1)[0]


class TorchCylinders:
    """
    A Pytorch representation of a batch of M cylinders (i.e. B elements in the batch,
    M cylinders per element). Any of these cylinders can have zero volume (these
    will be masked out during calculation of the various functions in this
    class, such as sdf).
    """

    def __init__(
        self,
        centers: torch.Tensor,
        radii: torch.Tensor,
        heights: torch.Tensor,
        quaternions: torch.Tensor,
    ):
        """
        :param centers torch.Tensor: Has dim [B, M, 3]
        :param radii torch.Tensor: Has dim [B, M, 1]
        :param heights torch.Tensor: Has dim [B, M, 1]
        :param quaternions torch.Tensor: Has dim [B, M, 4] with quaternions formatted as
                                   (w, x, y, z)
        """
        assert centers.ndim == 3
        assert radii.ndim == 3
        assert heights.ndim == 3
        assert quaternions.ndim == 3

        self.radii = radii
        self.heights = heights
        self.centers = centers

        # It's helpful to ensure the quaternions are normalized
        self.quats = quaternions / torch.linalg.norm(quaternions, dim=2)[:, :, None]
        self._init_frames()
        # Mask for nonzero volumes
        self.mask = ~torch.logical_or(
            torch.isclose(self.radii, torch.zeros(1).type_as(centers)).squeeze(-1),
            torch.isclose(self.heights, torch.zeros(1).type_as(centers)).squeeze(-1),
        )

    def geometrout(self):
        """
        Helper method to convert this into geometrout primitives
        """
        B, M, _ = self.centers.shape
        return [
            [
                Cylinder(
                    center=self.centers[bidx, midx, :].detach().cpu().numpy(),
                    radius=self.radii[bidx, midx, 0].detach().cpu().numpy(),
                    height=self.heights[bidx, midx, 0].detach().cpu().numpy(),
                    quaternion=self.quats[bidx, midx, :].detach().cpu().numpy(),
                )
                for midx in range(M)
                if self.mask[bidx, midx]
            ]
            for bidx in range(B)
        ]

    def _init_frames(self):
        """
        In order to calculate the SDF, we need to calculate the inverse
        transformation of the cylinder. This is because we are transforming
        points in the world frame into the cylinder frame.
        """
        # Initialize the inverse rotation
        w = self.quats[:, :, 0]
        x = -self.quats[:, :, 1]
        y = -self.quats[:, :, 2]
        z = -self.quats[:, :, 3]

        # Naming is a little disingenuous here because I'm multiplying everything by two,
        # but can't put numbers at the beginning of variable names.
        xx = 2 * torch.pow(x, 2)
        yy = 2 * torch.pow(y, 2)
        zz = 2 * torch.pow(z, 2)

        wx = 2 * w * x
        wy = 2 * w * y
        wz = 2 * w * z
        xy = 2 * x * y
        xz = 2 * x * z
        yz = 2 * y * z

        B, M, _ = self.centers.shape
        B = self.centers.size(0)
        M = self.centers.size(1)
        self.inv_frames = torch.zeros((B, M, 4, 4)).type_as(self.centers)
        self.inv_frames[:, :, 3, 3] = 1

        R = torch.stack(
            [
                torch.stack([1 - yy - zz, xy - wz, xz + wy], dim=2),
                torch.stack([xy + wz, 1 - xx - zz, yz - wx], dim=2),
                torch.stack([xz - wy, yz - wx, 1 - xx - yy], dim=2),
            ],
            dim=2,
        )
        Rt = torch.matmul(R, -1 * self.centers.unsqueeze(3)).squeeze(3)

        # Fill in the rotation matrices
        self.inv_frames[:, :, :3, :3] = R

        # Invert the transform by multiplying the inverse translation by the inverse rotation
        self.inv_frames[:, :, :3, 3] = Rt

    def sdf(self, points: torch.Tensor) -> torch.Tensor:
        """
        :param points torch.Tensor: The points with which to calculate the SDF, has
                                    dim [B, N, 3] (N is the number of points)
        :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF
                             value for each of the M cylinders), has dim [B, N]
        """
        assert points.ndim == 3
        assert points.size(0) == self.centers.size(0)
        if torch.all(~self.mask):
            return float("inf") * torch.ones(points.size(0), points.size(1)).type_as(
                points
            )

        homog_points = torch.cat(
            (
                points,
                torch.ones((points.size(0), points.size(1), 1)).type_as(points),
            ),
            dim=2,
        )
        # Next, project these points into their respective cylinder frames
        # Will return [B x M x N x 3]

        points_proj = torch.matmul(
            self.inv_frames[:, :, None, :, :], homog_points[:, None, :, :, None]
        ).squeeze(-1)[:, :, :, :3]
        B, M, N, _ = points_proj.shape
        masked_points = points_proj[self.mask]

        surface_distance_xy = torch.linalg.norm(masked_points[:, :, :2], dim=2)
        z_distance = masked_points[:, :, 2]

        half_extents_2d = torch.stack(
            (self.radii[self.mask], self.heights[self.mask] / 2), dim=2
        )
        points_2d = torch.stack((surface_distance_xy, z_distance), dim=2)
        distances_2d = torch.abs(points_2d) - half_extents_2d

        # This is distance only for points outside the box, all points inside return zero
        outside = torch.linalg.norm(
            torch.maximum(distances_2d, torch.zeros_like(distances_2d)), dim=2
        )
        # This is distance for points inside the box, all others return zero
        inner_max_distance_2d = torch.max(distances_2d, dim=2).values
        inside = torch.minimum(
            inner_max_distance_2d, torch.zeros_like(inner_max_distance_2d)
        )

        all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points)
        all_sdfs[self.mask] = outside + inside
        return torch.min(all_sdfs, dim=1)[0]

    def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor:
        """
        Calculates SDF values for a time sequence of point clouds
        :param points torch.Tensor: The batched sequence of point clouds with
                                    dimension [B, T, N, 3] (T in sequence length,
                                    N is number of points)
        :rtype torch.Tensor: The scene SDF for each point at each timestep
                             (i.e. the minimum SDF value across the M cylinders at
                             each timestep), has dim [B, T, N]
        """
        assert points.ndim == 4

        # We are going to need to map points in the global frame to the cylinder frame
        # First take the points and make them homogeneous by adding a one to the end

        assert points.size(0) == self.centers.size(0)
        if torch.all(~self.mask):
            return float("inf") * torch.ones(points.shape[:-1]).type_as(points)

        homog_points = torch.cat(
            (
                points,
                torch.ones((*points.shape[:-1], 1)).type_as(points),
            ),
            dim=3,
        )
        # Next, project these points into their respective cylinder frames
        # Will return [B x M x N x 3]

        points_proj = torch.matmul(
            self.inv_frames[:, :, None, None, :, :],
            homog_points[:, None, :, :, :, None],
        ).squeeze(-1)[:, :, :, :, :3]
        B, M, T, N, _ = points_proj.shape
        assert T == points.size(1)
        assert N == points.size(2)
        masked_points = points_proj[self.mask]

        surface_distance_xy = torch.linalg.norm(masked_points[:, :, :, :2], dim=-1)
        z_distance = masked_points[:, :, :, 2]

        half_extents_2d = torch.stack(
            (self.radii[self.mask], self.heights[self.mask] / 2), dim=2
        )[:, :, None, :]
        points_2d = torch.stack((surface_distance_xy, z_distance), dim=3)
        distances_2d = torch.abs(points_2d) - half_extents_2d

        # This is distance only for points outside the box, all points inside return zero
        outside = torch.linalg.norm(
            torch.maximum(distances_2d, torch.zeros_like(distances_2d)), dim=3
        )
        # This is distance for points inside the box, all others return zero
        inner_max_distance_2d = torch.max(distances_2d, dim=3).values
        inside = torch.minimum(
            inner_max_distance_2d, torch.zeros_like(inner_max_distance_2d)
        )

        all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points)
        all_sdfs[self.mask] = outside + inside
        return torch.min(all_sdfs, dim=1)[0]


def construct_mixed_point_cloud(
    obstacles: Sequence[Union[Sphere, Cuboid, Cylinder]], num_points: int
) -> np.ndarray:
    """
    Creates a random point cloud from a collection of obstacles. The points in
    the point cloud should be fairly(-ish) distributed amongst the obstacles based
    on their surface area.

    :param obstacles Sequence[Union[Sphere, Cuboid, Cylinder]]: The obstacles in the scene
    :param num_points int: The total number of points in the samples scene (not
                           the number of points per obstacle)
    :rtype np.ndarray: Has dim [N, 3] where N is num_points
    """
    point_set = []
    total_obstacles = len(obstacles)
    if total_obstacles == 0:
        return np.array([[]])

    # Allocate points based on obstacle surface area for even sampling
    surface_areas = np.array([o.surface_area for o in obstacles])
    total_area = np.sum(surface_areas)
    proportions = (surface_areas / total_area).tolist()

    indices = list(range(1, total_obstacles + 1))
    random.shuffle(indices)
    idx = 0

    for o, prop in zip(obstacles, proportions):
        sample_number = int(prop * num_points) + 500
        samples = o.sample_surface(sample_number)
        _points = indices[idx] * np.ones((sample_number, 4))
        _points[:, :3] = samples
        point_set.append(_points)
        idx += 1
    points = np.concatenate(point_set, axis=0)

    # Downsample to the desired number of points
    return points[np.random.choice(points.shape[0], num_points, replace=False), :]

mpinets/geometry.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

from typing import Tuple
from mpinets import utils
from mpinets.geometry import TorchCuboids, TorchCylinders
import torch.nn.functional as F
import torch
from robofin.pointcloud.torch import FrankaSampler


def point_match_loss(input_pc: torch.Tensor, target_pc: torch.Tensor) -> torch.Tensor:
    """
    A combination L1 and L2 loss to penalize large and small deviations between
    two point clouds

    :param input_pc torch.Tensor: Point cloud sampled from the network's output.
                                  Has dim [B, N, 3]
    :param target_pc torch.Tensor: Point cloud sampled from the supervision
                                   Has dim [B, N, 3]
    :rtype torch.Tensor: The single loss value
    """
    return F.mse_loss(input_pc, target_pc, reduction="mean") + F.l1_loss(
        input_pc, target_pc, reduction="mean"
    )


def collision_loss(
    input_pc: torch.Tensor,
    cuboid_centers: torch.Tensor,
    cuboid_dims: torch.Tensor,
    cuboid_quaternions: torch.Tensor,
    cylinder_centers: torch.Tensor,
    cylinder_radii: torch.Tensor,
    cylinder_heights: torch.Tensor,
    cylinder_quaternions: torch.Tensor,
) -> torch.Tensor:
    """
    Calculates the hinge loss, calculating whether the robot (represented as a
    point cloud) is in collision with any obstacles in the scene. Collision
    here actually means within 3cm of the obstacle--this is to provide stronger
    gradient signal to encourage the robot to move out of the way. Also, some of the
    primitives can have zero volume (i.e. a dim is zero for cuboids or radius or height is zero for cylinders).
    If these are zero volume, they will have infinite sdf values (and therefore be ignored by the loss).

    :param input_pc torch.Tensor: Points sampled from the robot's surface after it
                                  is placed at the network's output prediction. Has dim [B, N, 3]
    :param cuboid_centers torch.Tensor: Has dim [B, M1, 3]
    :param cuboid_dims torch.Tensor: Has dim [B, M1, 3]
    :param cuboid_quaternions torch.Tensor: Has dim [B, M1, 4]. Quaternion is formatted as w, x, y, z.
    :param cylinder_centers torch.Tensor: Has dim [B, M2, 3]
    :param cylinder_radii torch.Tensor: Has dim [B, M2, 1]
    :param cylinder_heights torch.Tensor: Has dim [B, M2, 1]
    :param cylinder_quaternions torch.Tensor: Has dim [B, M2, 4]. Quaternion is formatted as w, x, y, z.
    :rtype torch.Tensor: Returns the loss value aggregated over the batch
    """

    cuboids = TorchCuboids(
        cuboid_centers,
        cuboid_dims,
        cuboid_quaternions,
    )
    cylinders = TorchCylinders(
        cylinder_centers,
        cylinder_radii,
        cylinder_heights,
        cylinder_quaternions,
    )
    sdf_values = torch.minimum(cuboids.sdf(input_pc), cylinders.sdf(input_pc))
    return F.hinge_embedding_loss(
        sdf_values,
        -torch.ones_like(sdf_values),
        margin=0.03,
        reduction="mean",
    )


class CollisionAndBCLossContainer:
    """
    A container class to hold the various losses. This is structured as a
    container because that allows it to cache the robot pointcloud sampler
    object. By caching this, we reduce parsing time when processing the URDF
    and allow for a consistent random pointcloud (consistent per-GPU, that is)
    """

    def __init__(
        self,
    ):
        self.fk_sampler = None
        self.num_points = 1024

    def __call__(
        self,
        input_normalized: torch.Tensor,
        cuboid_centers: torch.Tensor,
        cuboid_dims: torch.Tensor,
        cuboid_quaternions: torch.Tensor,
        cylinder_centers: torch.Tensor,
        cylinder_radii: torch.Tensor,
        cylinder_heights: torch.Tensor,
        cylinder_quaternions: torch.Tensor,
        target_normalized: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        This method calculates both constituent loss function after loading,
        and then caching, a fixed robot point cloud sampler (i.e. the task
        spaces sampled are always the same, as opposed to a random point cloud).
        The fixed point cloud is important for loss calculation so that
        it's possible to take mse between the two pointclouds.

        :param input_normalized torch.Tensor: Has dim [B, 7] and is always between -1 and 1
        :param cuboid_centers torch.Tensor: Has dim [B, M1, 3]
        :param cuboid_dims torch.Tensor: Has dim [B, M1, 3]
        :param cuboid_quaternions torch.Tensor: Has dim [B, M1, 4]. Quaternion is formatted as w, x, y, z.
        :param cylinder_centers torch.Tensor: Has dim [B, M2, 3]
        :param cylinder_radii torch.Tensor: Has dim [B, M2, 1]
        :param cylinder_heights torch.Tensor: Has dim [B, M2, 1]
        :param cylinder_quaternions torch.Tensor: Has dim [B, M2, 4]. Quaternion is formatted as w, x, y, z.
        :param target_normalized torch.Tensor: Has dim [B, 7] and is always between -1 and 1
        :rtype Tuple[torch.Tensor, torch.Tensor]: The two losses aggregated over the batch
        """
        if self.fk_sampler is None:
            self.fk_sampler = FrankaSampler(
                input_normalized.device,
                num_fixed_points=self.num_points,
                use_cache=True,
                with_base_link=False,  # Remove base link because this isn't controllable anyway
            )
        input_pc = self.fk_sampler.sample(
            utils.unnormalize_franka_joints(input_normalized),
        )
        target_pc = self.fk_sampler.sample(
            utils.unnormalize_franka_joints(target_normalized),
        )
        return (
            collision_loss(
                input_pc,
                cuboid_centers,
                cuboid_dims,
                cuboid_quaternions,
                cylinder_centers,
                cylinder_radii,
                cylinder_heights,
                cylinder_quaternions,
            ),
            point_match_loss(input_pc, target_pc),
        )
mpinets/loss.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import numpy as np

from robofin.robots import FrankaRobot, FrankaGripper, FrankaRealRobot
from robofin.bullet import Bullet
import time
import pickle
from pathlib import Path
import itertools
from robofin.collision import FrankaSelfCollisionChecker
from geometrout.transform import SE3, SO3
from termcolor import colored
import logging
from mpinets.types import Obstacles

from geometrout.primitive import Cuboid, Sphere, Cylinder
from typing import Sequence, Union, List, Tuple, Any, Dict, Optional
from mpinets.types import Obstacles, Trajectory
from mpinets.third_party.sparc import sparc

try:
    import lula
except ImportError:
    logging.info("Lula not available. Falling back to bullet only")
    USE_LULA = False
else:
    USE_LULA = True


def percent_true(arr: Sequence) -> float:
    """
    Returns the percent true of a boolean sequence or the percent nonzero of a numerical sequence

    :param arr Sequence: The input sequence
    :rtype float: The percent
    """
    return 100 * np.count_nonzero(arr) / len(arr)


class Evaluator:
    """
    This class can be used to evaluate a whole set of environments and data
    """

    def __init__(self, fabric_urdf_path: str = None, gui: bool = False):
        """
        Initializes the evaluator class

        :param gui bool: Whether to visualize trajectories (and showing visually whether they are failures)
        """
        self.sim = Bullet(gui=False)
        self.selfcc = FrankaSelfCollisionChecker()
        self.hd_sim_robot = self.sim.load_robot(
            FrankaRobot, hd=True, default_prismatic_value=0.025
        )
        self.ld_sim_robot = self.sim.load_robot(
            FrankaRobot, hd=False, default_prismatic_value=0.025
        )

        self.self_collision_sim = Bullet(gui=False)
        self.self_collision_robot = self.self_collision_sim.load_robot(
            FrankaRobot, hd=True, default_prismatic_value=0.025
        )
        self.groups = {}
        self.current_group = None
        self.current_group_key = None
        if fabric_urdf_path is None:
            self.fabric_urdf_path = "/scratch/jayaram.reddy/issac_sim/library/isaac_sim-2022.2.1/exts/omni.isaac.motion_generation/motion_policy_configs/franka/lula_franka_gen.urdf"
        else:
            self.fabric_urdf_path = fabric_urdf_path
        if gui:
            self.gui_sim = Bullet(gui=True)
            self.gui_robot = self.gui_sim.load_robot(
                FrankaRobot, hd=True, default_prismatic_value=0.025
            )
            self.gui_gripper = self.gui_sim.load_robot(FrankaGripper)
        else:
            self.gui_sim = None
            self.gui_robot = None
            self.gui_gripper = None

    def run_visuals(
        self,
        trajectory: Trajectory,
        dt: float,
        target: SE3,
        obstacles: Obstacles,
        success: bool,
    ):
        """
        Visualizes a trajectory and changes the color based on whether its a success or failure

        :param trajectory Trajectory: The trajectory to visualize
        :param dt float: The approximate timestep to use when visualizing
        :param target SE3: The target pose (in right gripper frame)
        :param obstacles Obstacles: The obstacles to visualize
        :param success bool: Whether the trajectory was a success
        """
        self.gui_sim.clear_all_obstacles()
        if success:
            obstacle_color = [181 / 255, 132 / 255, 1, 1]
        else:
            obstacle_color = [1, 0.51, 0, 1]
        collision_ids = self.gui_sim.load_primitives(
            # obstacles, [16 / 255, 6 / 255, 54 / 255, 1]
            obstacles,
            obstacle_color,
        )
        self.gui_gripper.marionette(target)
        self.gui_sim.load_cuboid(
            Cuboid(
                center=target.xyz, dims=[0.01, 0.01, 0.01], quaternion=target.so3.wxyz
            ),
            [16 / 255, 6 / 255, 54 / 255, 1],
        )
        for i, q in enumerate(trajectory):
            self.gui_robot.marionette(q)
            collision_points = self.gui_robot.get_collision_points(
                collision_ids, check_self=False
            )
            if len(collision_points):
                self.gui_sim.load_primitives(
                    [Sphere(center=c, radius=0.02) for c in collision_points],
                    [1, 0, 0, 0.5],
                )
                time.sleep(0.25)
            time.sleep(dt)

    def create_new_group(self, key: str):
        """
        Creates a new metric group (for a new setting, for example)

        :param key str: The key for this metric group
        """
        self.groups[key] = {}
        self.current_group_key = key
        self.current_group = self.groups[key]

    if USE_LULA:

        def _get_fabric(self, obstacles: Obstacles) -> lula.Fabric:
            urdf_path = self.fabric_urdf_path
            assert Path(urdf_path).exists()
            fabric_robot_description_path = str(
                Path(__file__).resolve().parent.parent
                / "config"
                / "franka_robot_description.yaml"
            )
            assert Path(fabric_robot_description_path).exists()
            fabric_config_path = str(
                Path(__file__).resolve().parent.parent
                / "config"
                / "franka_fabric_config.yaml"
            )
            assert Path(fabric_config_path).exists()

            # Load robot description.
            robot_description = lula.load_robot(
                fabric_robot_description_path, urdf_path
            )
            fabric_state = lula.create_fabric_state()

            world = lula.create_world()
            for o in obstacles:
                if o.is_zero_volume():
                    continue
                if isinstance(o, Cuboid):
                    box_obstacle_pose = lula.Pose3(o.pose.matrix)
                    box = lula.create_obstacle(lula.Obstacle.Type.CUBE)
                    box.set_attribute(
                        lula.Obstacle.Attribute.SIDE_LENGTHS, np.asarray(o.dims)
                    )
                    world.add_obstacle(box, box_obstacle_pose)
                elif isinstance(o, Cylinder):
                    cylinder_obstacle_pose = lula.Pose3(o.pose.matrix)
                    cylinder = lula.create_obstacle(lula.Obstacle.Type.CYLINDER)
                    cylinder.set_attribute(lula.Obstacle.Attribute.RADIUS, o.radius)
                    cylinder.set_attribute(lula.Obstacle.Attribute.HEIGHT, o.height)
                    world.add_obstacle(cylinder, cylinder_obstacle_pose)

            world_view = world.add_world_view()

            # Create RMPflow configuration.
            fabric_config = lula.create_fabric_config(
                fabric_config_path,
                robot_description,
                "right_gripper",
                world_view,
            )
            # Create RMPflow policy
            fabric = lula.create_fabric(fabric_config)
            return fabric

        def in_collision(self, trajectory: Trajectory, obstacles: Obstacles) -> bool:
            """
            Checks whether the trajectory is in collision according to all including
            collision checkers (using AND between different methods)

            :param trajectory Trajectory: The trajectory
            :param obstacles Obstacles: Obstacles to check
            :rtype bool: Whether there is a collision
            """
            fabric = self._get_fabric(obstacles)
            for i, q in enumerate(trajectory):
                self.hd_sim_robot.marionette(q)
                self.ld_sim_robot.marionette(q)
                colcheck1 = self.sim.in_collision(self.ld_sim_robot, check_self=True)   #self collision?
                colcheck2 = self.sim.in_collision(self.hd_sim_robot, check_self=True)
                colcheck3 = fabric.in_collision_with_obstacle(q.astype(np.double))
                print("-"*20)
                print(f"Collision Checks: 1: {colcheck1}\t2: {colcheck2}\t3: {colcheck3}")
                print("-"*20)
                if (
                    colcheck1
                    and colcheck2
                    and colcheck3
                ):
                    return True
                # if (
                #     self.sim.in_collision(self.ld_sim_robot, check_self=True)
                #     and self.sim.in_collision(self.hd_sim_robot, check_self=True)
                #     and fabric.in_collision_with_obstacle(q.astype(np.double))
                # ):
                #     return True
            return False

    else:

        def in_collision(self, trajectory: Trajectory, obstacles: Obstacles) -> bool:
            """
            Checks whether the trajectory is in collision according to all including
            collision checkers (using AND between different methods)

            :param trajectory Trajectory: The trajectory
            :param obstacles Obstacles: Obstacles to check
            :rtype bool: Whether there is a collision
            """
            for i, q in enumerate(trajectory):
                self.hd_sim_robot.marionette(q)
                self.ld_sim_robot.marionette(q)
                # if self.sim.in_collision(
                #     self.ld_sim_robot, check_self=True
                # ) and self.sim.in_collision(self.hd_sim_robot, check_self=True):
                #     return True
                colcheck1 = self.sim.in_collision(self.ld_sim_robot, check_self=True)   #self collision?
                colcheck2 = self.sim.in_collision(self.hd_sim_robot, check_self=True)
                print("-"*20)
                print(f"Collision Checks: 1: {colcheck1}\t2: {colcheck2}")
                print("-"*20)
                if (
                    colcheck1
                    and colcheck2
                ):
                    return True
            return False

    def has_self_collision(self, trajectory: Trajectory) -> bool:
        """
        Checks whether there is a self collision (using OR between different methods)

        :param trajectory Trajectory: The trajectory
        :rtype bool: Whether there is a self collision
        """
        for i, q in enumerate(trajectory):
            self.self_collision_robot.marionette(q)
            if self.self_collision_sim.in_collision(
                self.self_collision_robot, check_self=True
            ) or self.selfcc.has_self_collision(q):
                return True
        return False

    def in_collision(self, trajectory: Trajectory, obstacles: Obstacles) -> bool:
        """
        Checks whether the trajectory is in collision according to all including
        collision checkers (using AND between different methods)

        :param trajectory Trajectory: The trajectory
        :param obstacles Obstacles: Obstacles to check
        :rtype bool: Whether there is a collision
        """
        if USE_LULA:
            fabric = self._get_fabric(obstacles)
        for i, q in enumerate(trajectory):
            self.hd_sim_robot.marionette(q)
            self.ld_sim_robot.marionette(q)
            if self.sim.in_collision(
                self.ld_sim_robot, check_self=True
            ) and self.sim.in_collision(self.hd_sim_robot, check_self=True):
                if not USE_LULA:
                    return True
                if fabric.in_collision_with_obstacle(q.astype(np.double)):
                    return True
        return False

    def get_collision_depths(
        self, trajectory: Trajectory, obstacles: Obstacles
    ) -> List[float]:
        """
        Get all the collision depths for a trajectory (sometimes can report strange
        values due to inconsistent Bullet collision checking)

        :param trajectory Trajectory: The trajectory
        :param obstacles Obstacles: The obstacles
        :rtype List[float]: A list of all collision depths
        """
        all_depths = []
        for i, q in enumerate(trajectory):
            self.hd_sim_robot.marionette(q)
            depths = self.hd_sim_robot.get_collision_depths(self.sim.obstacle_ids)
            all_depths.extend(depths)
        return all_depths

    @staticmethod
    def violates_joint_limits(trajectory: Trajectory) -> bool:
        """
        Checks whether any configuration in the trajectory violates joint limits

        :param trajectory Trajectory: The trajectory
        :rtype bool: Whether there is a joint limit violation
        """
        for i, q in enumerate(trajectory):
            if not FrankaRobot.within_limits(q):
                return True
        return False

    def has_physical_violation(
        self, trajectory: Trajectory, obstacles: Obstacles
    ) -> bool:
        """
        Checks whether there is any physical violation (collision, self collision, joint limit violation)

        :param trajectory Trajectory: The trajectory
        :param obstacles Obstacles: The obstacles in the scene
        :rtype bool: Whether there is at least one physical violation
        """
        return (
            self.in_collision(trajectory, obstacles)
            or self.violates_joint_limits(trajectory)
            or self.has_self_collision(trajectory)
        )

    @staticmethod
    def check_final_position(final_pose: SE3, target: SE3) -> float:
        """
        Gets the number of centimeters between the final pose and target

        :param final_pose SE3: The final pose of the trajectory
        :param target SE3: The target pose
        :rtype float: The distance in centimeters
        """
        return 100 * np.linalg.norm(final_pose._xyz - target._xyz)

    @staticmethod
    def check_final_orientation(final_orientation: SO3, target: SO3) -> float:
        """
        Gets the number of degrees between the final orientation and the target orientation

        :param final_orientation SO3: The final orientation
        :param target SO3: The final target orientation
        :rtype float: The rotational distance in degrees
        """
        return np.abs(
            (final_orientation._quat * target._quat.conjugate).radians * 180 / np.pi
        )

    @staticmethod
    def check_final_region(
        final_pose: SE3,
        target_volume: Union[Cuboid, Cylinder, Sphere],
        negative_volumes: Obstacles,
    ) -> bool:
        """
        Checks that the final pose is in the correct region if one is specified. For example,
        when reaching inside a drawer, it is not sufficient to be close to the final target as
        a condition of success. Instead, the final pose might be in the correct drawer
        (and likewise not in the incorrect drawer)

        :param final_pose SE3: The final pose of the trajectory
        :param target_volume Union[Cuboid, Cylinder, Sphere]: The target volume to be inside
        :param negative_volumes Obstacles: Volumes to be outside of
        :rtype bool: Whether the check is successful and the target is in the right region
                     and outside the wrong regions
        """
        return target_volume.sdf(final_pose._xyz) <= 0 and np.all(
            [v.sdf(final_pose._xyz) > 0 for v in negative_volumes]
        )

    @staticmethod
    def calculate_smoothness(trajectory: Trajectory, dt: float) -> Tuple[float, float]:
        """
        Calculate trajectory smoothness using SPARC

        :param trajectory Trajectory: The trajectory
        :param dt float: The timestep in between consecutive steps of the trajectory
        :rtype Tuple[float, float]: The SPARC in configuration space and end effector space
        """
        configs = np.asarray(trajectory)
        assert configs.ndim == 2 and configs.shape[1] == 7
        config_movement = np.linalg.norm(np.diff(configs, 1, axis=0) / dt, axis=1)
        assert len(config_movement) == len(configs) - 1
        config_sparc, _, _ = sparc(config_movement, 1.0 / dt)

        eff_positions = np.asarray(
            [FrankaRobot.fk(q, eff_frame="right_gripper")._xyz for q in trajectory]
        )
        assert eff_positions.ndim == 2 and eff_positions.shape[1] == 3
        eff_movement = np.linalg.norm(np.diff(eff_positions, 1, axis=0) / dt, axis=1)
        assert len(eff_movement) == len(eff_positions) - 1
        eff_sparc, _, _ = sparc(eff_movement, 1.0 / dt)

        return config_sparc, eff_sparc

    def calculate_eff_path_lengths(self, trajectory: Trajectory) -> Tuple[float, float]:
        """
        Calculate the end effector path lengths (position and orientation).
        Orientation is in degrees.

        :param trajectory Trajectory: The trajectory
        :rtype Tuple[float, float]: The path lengths (position, orientation)
        """
        eff_poses = [FrankaRobot.fk(q, eff_frame="right_gripper") for q in trajectory]

        eff_positions = np.asarray([pose._xyz for pose in eff_poses])
        assert eff_positions.ndim == 2 and eff_positions.shape[1] == 3
        position_step_lengths = np.linalg.norm(
            np.diff(eff_positions, 1, axis=0), axis=1
        )
        eff_position_path_length = sum(position_step_lengths)

        eff_quaternions = [pose.so3._quat for pose in eff_poses]
        eff_orientation_path_length = 0
        for qi, qj in zip(eff_quaternions[:-1], eff_quaternions[1:]):
            eff_orientation_path_length += np.abs(
                np.degrees((qj * qi.conjugate).radians)
            )
        return eff_position_path_length, eff_orientation_path_length

    def evaluate_trajectory(
        self,
        trajectory: Trajectory,
        dt: float,
        target: SE3,
        obstacles: Obstacles,
        target_volume: Union[Cuboid, Cylinder, Sphere],
        target_negative_volumes: Obstacles,
        time: float,
        skip_metrics: bool = False,
    ):
        """
        Evaluates a single trajectory and stores the metrics in the current group.
        Will visualize and print relevant info if `self.gui` is `True`

        :param trajectory Trajectory: The trajectory
        :param dt float: The time step for the trajectory
        :param target SE3: The target pose
        :param obstacles Obstacles: The obstacles in the scene
        :param target_volume Union[Cuboid, Cylinder, Sphere]: The target volume for the trajectory
        :param target_negative_volumes Obstacles: Volumes that the target should definitely be outside
        :param time float: The time taken to calculate the trajectory
        :param skip_metrics bool: Whether to skip the path metrics (for example if it's a feasibility planner that failed)
        """

        def add_metric(key, value):
            self.current_group[key] = self.current_group.get(key, []) + [value]

        if skip_metrics:
            add_metric("success", False)
            add_metric("time", np.inf)
            add_metric("skips", True)
            return
        self.sim.clear_all_obstacles()
        self.sim.load_primitives(obstacles)

        in_collision = self.in_collision(trajectory, obstacles)
        if in_collision:
            collision_depths = (
                self.get_collision_depths(trajectory, obstacles) if in_collision else []
            )
        else:
            collision_depths = []
        add_metric("collision_depths", collision_depths)
        add_metric("collision", in_collision)
        add_metric("joint_limit_violation", self.violates_joint_limits(trajectory))
        add_metric("self_collision", self.has_self_collision(trajectory))

        physical_violation = self.has_physical_violation(trajectory, obstacles)
        add_metric("physical_violations", physical_violation)

        final_pose = FrankaRobot.fk(trajectory[-1], eff_frame="right_gripper")

        position_error = self.check_final_position(final_pose, target)
        add_metric("position_error", position_error)

        orientation_error = self.check_final_orientation(final_pose.so3, target.so3)
        add_metric("orientation_error", orientation_error)

        config_smoothness, eff_smoothness = self.calculate_smoothness(trajectory, dt)
        add_metric("config_smoothness", config_smoothness)
        add_metric("eff_smoothness", eff_smoothness)

        (
            eff_position_path_length,
            eff_orientation_path_length,
        ) = self.calculate_eff_path_lengths(trajectory)
        add_metric("eff_position_path_length", eff_position_path_length)
        add_metric("eff_orientation_path_length", eff_orientation_path_length)

        # Sometimes the target is inside a negative volume. This is obviously a bad negative volume
        corrected_negative_volumes = [
            v for v in target_negative_volumes if v.sdf(target._xyz) > 0
        ]
        correct_final_region = self.check_final_region(
            final_pose, target_volume, corrected_negative_volumes
        )

        success = (
            position_error < 1
            and correct_final_region
            and orientation_error < 15
            and not physical_violation
        )

        add_metric("success", success)
        add_metric("time", time)
        add_metric("num_steps", len(trajectory))

        def msg(metric, is_error):
            return colored(metric, "red") if is_error else colored(metric, "green")

        if self.gui_sim is not None and len(collision_depths) > 0:
            print("Position Error:", msg(position_error, position_error > 1))
            print("Orientation Error:", msg(orientation_error, orientation_error > 15))
            print("Config SPARC:", msg(config_smoothness, config_smoothness > -1.6))
            print("End Eff SPARC:", msg(eff_smoothness, eff_smoothness > -1.6))
            print(f"End Eff Position Path Length: {eff_position_path_length}")
            print(f"End Eff Orientation Path Length: {eff_orientation_path_length}")
            print("Physical violation:", msg(physical_violation, physical_violation))
            print(
                "Collisions:",
                msg(
                    self.in_collision(trajectory, obstacles),
                    self.in_collision(trajectory, obstacles),
                ),
            )
            if len(collision_depths) > 0:
                print(f"Mean Collision Depths: {100 * np.max(collision_depths)}")
                print(f"Median Collision Depths: {100 * np.median(collision_depths)}")
            print(
                "Joint Limit Violation:",
                msg(
                    self.violates_joint_limits(trajectory),
                    self.violates_joint_limits(trajectory),
                ),
            )
            print(
                "Self collision:",
                msg(
                    self.has_self_collision(trajectory),
                    self.has_self_collision(trajectory),
                ),
            )
            print(
                "Correct Region:", msg(correct_final_region, not correct_final_region)
            )
            self.run_visuals(trajectory, dt, target, obstacles, success)

    @staticmethod
    def metrics(group: Dict[str, Any]) -> Dict[str, float]:
        """
        Calculates the metrics for a specific group

        :param group Dict[str, Any]: The group of results
        :rtype Dict[str, float]: The metrics
        """
        success = percent_true(group["success"])
        one_cm = percent_true(np.array(group["position_error"]) < 1)
        five_cm = percent_true(np.array(group["position_error"]) < 5)
        fifteen_deg = percent_true(np.array(group["orientation_error"]) < 15)
        thirty_deg = percent_true(np.array(group["orientation_error"]) < 30)
        turn_around = percent_true(np.array(group["orientation_error"]) > 165)

        physical = percent_true(group["physical_violations"])
        config_smoothness = np.mean(group["config_smoothness"])
        eff_smoothness = np.mean(group["eff_smoothness"])
        all_eff_position_path_lengths = np.asarray(group["eff_position_path_length"])
        all_eff_orientation_path_lengths = np.asarray(
            group["eff_orientation_path_length"]
        )
        all_times = np.asarray(group["time"])

        is_smooth = percent_true(
            np.logical_and(
                np.asarray(group["config_smoothness"]) < -1.6,
                np.asarray(group["eff_smoothness"]) < -1.6,
            )
        )

        # Only use filtered successes for eff path length because we only
        # save the metric when it's an unskipped trajectory
        # TODO maybe clean this up so that it's not so if/else-y and either
        # all metrics are saved or only the ones that aren't skipped save
        skips = []
        if "skips" in group:
            # Needed when there are skips
            successes = np.asarray(group["success"])
            unskipped_successes = successes[~np.isinf(all_times)]
            skips = group["skips"]
        else:
            unskipped_successes = group["success"]

        success_position_path_lengths = all_eff_position_path_lengths[
            unskipped_successes
        ]
        success_orientation_path_lengths = all_eff_orientation_path_lengths[
            unskipped_successes
        ]
        eff_position_path_length = (
            np.mean(success_position_path_lengths),
            np.std(success_position_path_lengths),
        )
        eff_orientation_path_length = (
            np.mean(success_orientation_path_lengths),
            np.std(success_orientation_path_lengths),
        )
        success_times = all_times[group["success"]]
        time = (
            np.mean(success_times),
            np.std(success_times),
        )

        collision = percent_true(group["collision"])
        joint_limit = percent_true(group["joint_limit_violation"])
        self_collision = percent_true(group["self_collision"])
        depths = np.array(
            list(itertools.chain.from_iterable(group["collision_depths"]))
        )
        all_num_steps = np.asarray(group["num_steps"])
        success_num_steps = all_num_steps[unskipped_successes]
        step_time = (
            np.mean(success_times / success_num_steps),
            np.std(success_times / success_num_steps),
        )
        return {
            "success": success,
            "total": len(group["success"]),
            "skips": len(skips),
            "time": time,
            "step time": step_time,
            "env collision": collision,
            "self collision": self_collision,
            "joint violation": joint_limit,
            "physical violations": physical,
            "average collision depth": 100 * np.mean(depths),
            "median collision depth": 100 * np.median(depths),
            "1 cm": one_cm,
            "5 cm": five_cm,
            "15 deg": fifteen_deg,
            "30 deg": thirty_deg,
            "165 deg": turn_around,
            "is smooth": is_smooth,
            "average config sparc": config_smoothness,
            "average eff sparc": eff_smoothness,
            "eff position path length": eff_position_path_length,
            "eff orientation path length": eff_orientation_path_length,
        }

    @staticmethod
    def print_metrics(group: Dict[str, Any]):
        """
        Prints the metrics in an easy to read format

        :param group Dict[str, float]: The group of results
        """
        metrics = Evaluator.metrics(group)
        print(f"Total problems: {metrics['total']}")
        print(f"# Skips (Hard Failures): {metrics['skips']}")
        print(f"% Success: {metrics['success']:4.2f}")
        print(f"% Within 1cm: {metrics['1 cm']:4.2f}")
        print(f"% Within 5cm: {metrics['5 cm']:4.2f}")
        print(f"% Within 15deg: {metrics['15 deg']:4.2f}")
        print(f"% Within 30deg: {metrics['30 deg']:4.2f}")
        print(f"% Within 15deg of 180: {metrics['165 deg']:4.2f}")
        print(f"% With Environment Collision: {metrics['env collision']:4.2f}")
        print(f"% With Self Collision: {metrics['self collision']:4.2f}")
        print(f"% With Joint Limit Violations: {metrics['joint violation']:4.2f}")
        print(f"% With Self Collision: {metrics['self collision']:4.2f}")
        print(f"Average Collision Depth (cm): {metrics['average collision depth']}")
        print(f"Median Collision Depth (cm): {metrics['median collision depth']}")
        print(f"% With Physical Violations: {metrics['physical violations']:4.2f}")
        print(f"Average Config SPARC: {metrics['average config sparc']:4.2f}")
        print(f"Average End Eff SPARC: {metrics['average eff sparc']:4.2f}")
        print(f"% Smooth: {metrics['is smooth']:4.2f}")
        print(
            "Average End Eff Position Path Length:"
            f" {metrics['eff position path length'][0]:4.2f}"
            f" ± {metrics['eff position path length'][1]:4.2f}"
        )
        print(
            "Average End Eff Orientation Path Length:"
            f" {metrics['eff orientation path length'][0]:4.2f}"
            f" ± {metrics['eff orientation path length'][1]:4.2f}"
        )
        print(f"Average Time: {metrics['time'][0]:4.2f} ± {metrics['time'][1]:4.2f}")
        print(
            "Average Time Per Step (Not Always Valuable):"
            f" {metrics['step time'][0]:4.6f}"
            f" ± {metrics['step time'][1]:4.6f}"
        )

    def save_group(self, directory: str, test_name: str, key: Optional[str] = None):
        """
        Save the results of a single group

        :param directory str: The directory in which to save the results
        :param test_name str: The name of this specific test
        :param key Optional[str]: The group key to use. If not specified will use the current group
        """
        if key is None:
            group = self.current_group
        else:
            group = self.groups[key]
        save_path = Path(directory) / f"{test_name}_{self.current_group_key}.pkl"
        print(f"Saving group metrics to {save_path}")
        with open(save_path, "wb") as f:
            pickle.dump(group, f)

    def save(self, directory: str, test_name: str):
        """
        Save all the groups

        :param directory str: The directory name in which to save results
        :param test_name str: The test name (used as the file name)
        """
        save_path = Path(directory) / f"{test_name}_metrics.pkl"
        print(f"Metrics will save to {save_path}")
        with open(save_path, "wb") as f:
            pickle.dump(self.groups, f)

    def print_group_metrics(self, key: Optional[str] = None):
        """
        Prints out the metrics for a specific group

        :param key Optional[str]: The group key (if none specified, will use current group)
        """
        if key is not None:
            self.current_group = self.groups[key]
            self.current_group_key = key
        return self.print_metrics(self.current_group)

    def print_overall_metrics(self):
        """
        Prints the metrics for the aggregated results over all groups
        """
        supergroup = {}
        keys = set()
        for group in self.groups.values():
            for key in group.keys():
                keys.add(key)

        for key in keys:
            metrics = []
            for group in self.groups.values():
                metrics.extend(group.get(key, []))
            supergroup[key] = metrics
        return self.print_metrics(supergroup)

mpinets/metrics.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import torch
from torch import nn
from robofin.pointcloud.torch import FrankaSampler, FrankaCollisionSampler
import pytorch_lightning as pl
from pointnet2_ops.pointnet2_modules import PointnetSAModule

from mpinets import loss
from mpinets.utils import unnormalize_franka_joints
from mpinets.geometry import TorchCuboids, TorchCylinders
from typing import List, Tuple, Sequence, Dict, Callable


class MotionPolicyNetwork(pl.LightningModule):
    """
    The architecture laid out here is the default architecture laid out in the
    Motion Policy Networks paper (Fishman, et. al, 2022).
    """

    def __init__(self):
        """
        Constructs the model
        """
        super().__init__()
        self.point_cloud_encoder = MPiNetsPointNet()
        self.feature_encoder = nn.Sequential(
            nn.Linear(7, 32),
            nn.LeakyReLU(),
            nn.Linear(32, 64),
            nn.LeakyReLU(),
            nn.Linear(64, 128),
            nn.LeakyReLU(),
            nn.Linear(128, 128),
            nn.LeakyReLU(),
            nn.Linear(128, 64),
        )
        self.decoder = nn.Sequential(
            nn.Linear(2048 + 64, 512),
            nn.LeakyReLU(),
            nn.Linear(512, 256),
            nn.LeakyReLU(),
            nn.Linear(256, 128),
            nn.LeakyReLU(),
            nn.Linear(128, 7),
        )

    def configure_optimizers(self):
        """
        A standard method in PyTorch lightning to set the optimizer
        """
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)
        return optimizer

    def forward(self, xyz: torch.Tensor, q: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        """
        Passes data through the network to produce an output

        :param xyz torch.Tensor: Tensor representing the point cloud. Should
                                      have dimensions of [B x N x 4] where B is the batch
                                      size, N is the number of points and 4 is because there
                                      are three geometric dimensions and a segmentation mask
        :param q torch.Tensor: The current robot configuration normalized to be between
                                    -1 and 1, according to each joint's range of motion
        :rtype torch.Tensor: The displacement to be applied to the current configuration to get
                     the position at the next step (still in normalized space)
        """
        pc_encoding = self.point_cloud_encoder(xyz)
        feature_encoding = self.feature_encoder(q)
        x = torch.cat((pc_encoding, feature_encoding), dim=1)
        return self.decoder(x)


class TrainingMotionPolicyNetwork(MotionPolicyNetwork):
    """
    An version of the MotionPolicyNetwork model that has additional attributes
    necessary during training (or using the validation step outside of the
    training process). This class is a valid model, but it's overkill when
    doing real robot inference and, for example, point cloud sampling is
    done by an outside process (such as downsampling point clouds from a point cloud).
    """

    def __init__(
        self,
        num_robot_points: int,
        point_match_loss_weight: float,
        collision_loss_weight: float,
    ):
        """
        Creates the network and assigns additional parameters for training


        :param num_robot_points int: The number of robot points used when resampling
                                     the robot points during rollouts (used in validation)
        :param point_match_loss_weight float: The weight assigned to the behavior
                                              cloning loss.
        :param collision_loss_weight float: The weight assigned to the collision loss
        :rtype Self: An instance of the network
        """
        super().__init__()
        self.num_robot_points = num_robot_points
        self.point_match_loss_weight = point_match_loss_weight
        self.collision_loss_weight = collision_loss_weight
        self.fk_sampler = None
        self.collision_sampler = None
        self.loss_fun = loss.CollisionAndBCLossContainer()

    def rollout(
        self,
        batch: Dict[str, torch.Tensor],
        rollout_length: int,
        sampler: Callable[[torch.Tensor], torch.Tensor],
        unnormalize: bool = False,
    ) -> List[torch.Tensor]:
        """
        Rolls out the policy an arbitrary length by calling it iteratively

        :param batch Dict[str, torch.Tensor]: A data batch coming from the
                                              data loader--should already be
                                              on the correct device
        :param rollout_length int: The number of steps to roll out (not including the start)
        :param sampler Callable[[torch.Tensor], torch.Tensor]: A function that takes a batch of robot
                                                               configurations [B x 7] and returns a batch of
                                                               point clouds samples on the surface of that robot
        :param unnormalize bool: Whether to return the whole trajectory unnormalized
                                 (i.e. converted back into joint space)
        :rtype List[torch.Tensor]: The entire trajectory batch, i.e. a list of
                                   configuration batches including the starting
                                   configurations where each element in the list
                                   corresponds to a timestep. For example, the
                                   first element of each batch in the list would
                                   be a single trajectory.
        """
        xyz, q = (
            batch["xyz"],
            batch["configuration"],
        )
        # This block is to adapt for the case where we only want to roll out a
        # single trajectory
        if q.ndim == 1:
            xyz = xyz.unsqueeze(0)
            q = q.unsqueeze(0)
        if unnormalize:
            q_unnorm = unnormalize_franka_joints(q)
            assert isinstance(q_unnorm, torch.Tensor)
            trajectory = [q_unnorm]
        else:
            trajectory = [q]

        for i in range(rollout_length):
            q = torch.clamp(q + self(xyz, q), min=-1, max=1)
            q_unnorm = unnormalize_franka_joints(q)
            assert isinstance(q_unnorm, torch.Tensor)
            q_unnorm = q_unnorm.type_as(q)
            if unnormalize:
                trajectory.append(q_unnorm)
            else:
                trajectory.append(q)

            samples = sampler(q_unnorm).type_as(xyz)
            xyz[:, : samples.shape[1], :3] = samples

        return trajectory

    def training_step(  # type: ignore[override]
        self, batch: Dict[str, torch.Tensor], batch_idx: int
    ) -> torch.Tensor:
        """
        A function called automatically by Pytorch Lightning during training.
        This function handles the forward pass, the loss calculation, and what to log

        :param batch Dict[str, torch.Tensor]: A data batch coming from the
                                                   data loader--should already be
                                                   on the correct device
        :param batch_idx int: The index of the batch (not used by this function)
        :rtype torch.Tensor: The overall weighted loss (used for backprop)
        """
        xyz, q = (
            batch["xyz"],
            batch["configuration"],
        )
        y_hat = torch.clamp(q + self(xyz, q), min=-1, max=1)
        (
            cuboid_centers,
            cuboid_dims,
            cuboid_quats,
            cylinder_centers,
            cylinder_radii,
            cylinder_heights,
            cylinder_quats,
            supervision,
        ) = (
            batch["cuboid_centers"],
            batch["cuboid_dims"],
            batch["cuboid_quats"],
            batch["cylinder_centers"],
            batch["cylinder_radii"],
            batch["cylinder_heights"],
            batch["cylinder_quats"],
            batch["supervision"],
        )
        collision_loss, point_match_loss = self.loss_fun(
            y_hat,
            cuboid_centers,
            cuboid_dims,
            cuboid_quats,
            cylinder_centers,
            cylinder_radii,
            cylinder_heights,
            cylinder_quats,
            supervision,
        )
        self.log("point_match_loss", point_match_loss)
        self.log("collision_loss", collision_loss)
        val_loss = (
            self.point_match_loss_weight * point_match_loss
            + self.collision_loss_weight * collision_loss
        )
        self.log("val_loss", val_loss)
        return val_loss

    def sample(self, q: torch.Tensor) -> torch.Tensor:
        """
        Samples a point cloud from the surface of all the robot's links

        :param q torch.Tensor: Batched configuration in joint space
        :rtype torch.Tensor: Batched point cloud of size [B, self.num_robot_points, 3]
        """
        assert self.fk_sampler is not None
        return self.fk_sampler.sample(q, self.num_robot_points)

    def validation_step(  # type: ignore[override]
        self, batch: Dict[str, torch.Tensor], batch_idx: int
    ) -> torch.Tensor:
        """
        This is a Pytorch Lightning function run automatically across devices
        during the validation loop

        :param batch Dict[str, torch.Tensor]: The batch coming from the dataloader
        :param batch_idx int: The index of the batch (not used by this function)
        :rtype torch.Tensor: The loss values which are to be collected into summary stats
        """

        # These are defined here because they need to be set on the correct devices.
        # The easiest way to do this is to do it at call-time
        if self.fk_sampler is None:
            self.fk_sampler = FrankaSampler(self.device, use_cache=True)
        if self.collision_sampler is None:
            self.collision_sampler = FrankaCollisionSampler(
                self.device, with_base_link=False
            )
        rollout = self.rollout(batch, 69, self.sample, unnormalize=True)

        assert self.fk_sampler is not None  # Necessary for mypy to type properly
        eff = self.fk_sampler.end_effector_pose(rollout[-1])
        position_error = torch.linalg.vector_norm(
            eff[:, :3, -1] - batch["target_position"], dim=1
        )
        avg_target_error = torch.mean(position_error)

        cuboids = TorchCuboids(
            batch["cuboid_centers"],
            batch["cuboid_dims"],
            batch["cuboid_quats"],
        )
        cylinders = TorchCylinders(
            batch["cylinder_centers"],
            batch["cylinder_radii"],
            batch["cylinder_heights"],
            batch["cylinder_quats"],
        )

        B = batch["cuboid_centers"].size(0)
        rollout = torch.stack(rollout, dim=1)
        # Here is some Pytorch broadcasting voodoo to calculate whether each
        # rollout has a collision or not (looking to calculate the collision rate)
        assert rollout.shape == (B, 70, 7)
        rollout = rollout.reshape(-1, 7)
        has_collision = torch.zeros(B, dtype=torch.bool, device=self.device)
        collision_spheres = self.collision_sampler.compute_spheres(rollout)
        for radius, spheres in collision_spheres:
            num_spheres = spheres.shape[-2]
            sphere_sequence = spheres.reshape((B, -1, num_spheres, 3))
            sdf_values = torch.minimum(
                cuboids.sdf_sequence(sphere_sequence),
                cylinders.sdf_sequence(sphere_sequence),
            )
            assert sdf_values.shape == (B, 70, num_spheres)
            radius_collisions = torch.any(
                sdf_values.reshape((sdf_values.size(0), -1)) <= radius, dim=-1
            )
            has_collision = torch.logical_or(radius_collisions, has_collision)

        avg_collision_rate = torch.count_nonzero(has_collision) / B
        return {
            "avg_target_error": avg_target_error,
            "avg_collision_rate": avg_collision_rate,
        }

    def validation_step_end(  # type: ignore[override]
        self, batch_parts: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Called by Pytorch Lightning at the end of each validation step to
        aggregate across devices

        :param batch_parts Dict[str, torch.Tensor]: The parts accumulated from all devices
        :rtype Dict[str, torch.Tensor]: The average values across the devices
        """
        return {
            "avg_target_error": torch.mean(batch_parts["avg_target_error"]),
            "avg_collision_rate": torch.mean(batch_parts["avg_collision_rate"]),
        }

    def validation_epoch_end(  # type: ignore[override]
        self, validation_step_outputs: Sequence[Dict[str, torch.Tensor]]
    ):
        """
        Pytorch lightning method that aggregates stats from the validation loop and logs

        :param validation_step_outputs Sequence[Dict[str, torch.Tensor]]: The outputs from each
                                                                      validation step
        """
        avg_target_error = torch.mean(
            torch.stack([x["avg_target_error"] for x in validation_step_outputs])
        )
        self.log("avg_target_error", avg_target_error)

        avg_collision_rate = torch.mean(
            torch.stack([x["avg_collision_rate"] for x in validation_step_outputs])
        )
        self.log("avg_collision_rate", avg_collision_rate)


class MPiNetsPointNet(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self._build_model()

    def _build_model(self):
        """
        Assembles the model design into a ModuleList
        """
        self.SA_modules = nn.ModuleList()
        self.SA_modules.append(
            PointnetSAModule(
                npoint=512,
                radius=0.05,
                nsample=128,
                mlp=[1, 64, 64, 64],
                bn=False,
            )
        )
        self.SA_modules.append(
            PointnetSAModule(
                npoint=128,
                radius=0.3,
                nsample=128,
                mlp=[64, 128, 128, 256],
                bn=False,
            )
        )
        self.SA_modules.append(PointnetSAModule(mlp=[256, 512, 512, 1024], bn=False))

        self.fc_layer = nn.Sequential(
            nn.Linear(1024, 4096),
            nn.GroupNorm(16, 4096),
            nn.LeakyReLU(inplace=True),
            nn.Linear(4096, 2048),
            nn.GroupNorm(16, 2048),
            nn.LeakyReLU(inplace=True),
            nn.Linear(2048, 2048),
        )

    @staticmethod
    def _break_up_pc(pc: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Breaks up the point cloud into the xyz coordinates and segmentation mask

        :param pc torch.Tensor: Tensor with shape [B, N, M] where M is larger than 3.
                                The first three dimensions along the last axis will be x, y, z
        :rtype Tuple[torch.Tensor, torch.Tensor]: Two tensors, one with just xyz
            and one with the corresponding features
        """
        xyz = pc[..., 0:3].contiguous()
        features = pc[..., 3:].transpose(1, 2).contiguous()
        return xyz, features

    def forward(self, point_cloud: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        """
        Forward pass of the network

        :param point_cloud torch.Tensor: Has dimensions (B, N, 4)
                                              B is the batch size
                                              N is the number of points
                                              4 is x, y, z, segmentation_mask
                                              This tensor must be on the GPU (CPU tensors not supported)
        :rtype torch.Tensor: The output from the network
        """
        assert point_cloud.size(2) == 4
        
        xyz, features = self._break_up_pc(point_cloud)
        # print(xyz.shape)
        # print(features.shape)
        for module in self.SA_modules:
            xyz, features = module(xyz, features)

        return self.fc_layer(features.squeeze(-1))
mpinets/model.py
vishal-2000/EDMP
# Copyright (c) 2015, Sivakumar Balasubramanian <siva82kb@gmail.com>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# Not a contribution
# Changes made by NVIDIA CORPORATION & AFFILIATES or University of Washington enabling SPARC computation
# or otherwise documented as NVIDIA-proprietary are not a contribution and subject
# to the following terms and conditions:
#
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.


# Note: original code can be found on Github at
# https://github.com/siva82kb/SPARC/blob/master/scripts/smoothness.py

import numpy as np


def sparc(movement, fs, padlevel=4, fc=10.0, amp_th=0.05):
    """
    Calculates the smoothness of the given speed profile using the modified
    spectral arc length metric.
    Parameters
    ----------
    movement : np.array
               The array containing the movement speed profile.
    fs       : float
               The sampling frequency of the data.
    padlevel : integer, optional
               Indicates the amount of zero padding to be done to the movement
               data for estimating the spectral arc length. [default = 4]
    fc       : float, optional
               The max. cut off frequency for calculating the spectral arc
               length metric. [default = 10.]
    amp_th   : float, optional
               The amplitude threshold to used for determing the cut off
               frequency upto which the spectral arc length is to be estimated.
               [default = 0.05]
    Returns
    -------
    sal      : float
               The spectral arc length estimate of the given movement's
               smoothness.
    (f, Mf)  : tuple of two np.arrays
               This is the frequency(f) and the magntiude spectrum(Mf) of the
               given movement data. This spectral is from 0. to fs/2.
    (f_sel, Mf_sel) : tuple of two np.arrays
                      This is the portion of the spectrum that is selected for
                      calculating the spectral arc length.
    Notes
    -----
    This is the modfieid spectral arc length metric, which has been tested only
    for discrete movements.

    Examples
    --------
    >>> t = np.arange(-1, 1, 0.01)
    >>> move = np.exp(-5*pow(t, 2))
    >>> sal, _, _ = sparc(move, fs=100.)
    >>> '%.5f' % sal
    '-1.41403'
    """
    if np.allclose(movement, 0):
        print("All movement was 0, returning 0")
        return 0, None, None
    # Number of zeros to be padded.
    nfft = int(pow(2, np.ceil(np.log2(len(movement))) + padlevel))

    # Frequency
    f = np.arange(0, fs, fs / nfft)
    # Normalized magnitude spectrum
    Mf = abs(np.fft.fft(movement, nfft))
    Mf = Mf / max(Mf)

    # Indices to choose only the spectrum within the given cut off frequency
    # Fc.
    # NOTE: This is a low pass filtering operation to get rid of high frequency
    # noise from affecting the next step (amplitude threshold based cut off for
    # arc length calculation).
    fc_inx = ((f <= fc) * 1).nonzero()
    f_sel = f[fc_inx]
    Mf_sel = Mf[fc_inx]

    # Choose the amplitude threshold based cut off frequency.
    # Index of the last point on the magnitude spectrum that is greater than
    # or equal to the amplitude threshold.
    inx = ((Mf_sel >= amp_th) * 1).nonzero()[0]
    fc_inx = range(inx[0], inx[-1] + 1)
    f_sel = f_sel[fc_inx]
    Mf_sel = Mf_sel[fc_inx]

    # Calculate arc length
    new_sal = -sum(
        np.sqrt(
            pow(np.diff(f_sel) / (f_sel[-1] - f_sel[0]), 2) + pow(np.diff(Mf_sel), 2)
        )
    )
    return new_sal, (f, Mf), (f_sel, Mf_sel)

mpinets/third_party/sparc.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

from dataclasses import dataclass, field
from typing import List, Union, Optional, Dict, Sequence
import numpy as np
from geometrout.transform import SE3
from geometrout.primitive import Cuboid, Cylinder, Sphere


Obstacles = List[Union[Cuboid, Cylinder, Sphere]]
Trajectory = Sequence[Union[Sequence, np.ndarray]]


@dataclass
class PlanningProblem:
    """
    Defines a common interface to describe planning problems
    """

    target: SE3  # The target in the `right_gripper` frame
    target_volume: Union[Cuboid, Cylinder]
    q0: np.ndarray  # The starting configuration
    obstacles: Optional[Obstacles] = None  # The obstacles in the scene
    obstacle_point_cloud: Optional[np.ndarray] = None
    target_negative_volumes: Obstacles = field(default_factory=lambda: [])


ProblemSet = Dict[str, Dict[str, List[PlanningProblem]]]

mpinets/types.py
vishal-2000/EDMP
# MIT License
#
# Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

from typing import Union, Tuple

import numpy as np
import torch
from robofin.robots import FrankaRobot, FrankaRealRobot


def _normalize_franka_joints_numpy(
    batch_trajectory: np.ndarray,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> np.ndarray:
    """
    Normalizes joint angles to be within a specified range according to the Franka's
    joint limits. This is the numpy version

    :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The new limits to map to
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype np.ndarray: An array with the same dimensions as the input
    """
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = robot.JOINT_LIMITS
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF)
        or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF)
        or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF)
    )
    normalized = (batch_trajectory - franka_limits[:, 0]) / (
        franka_limits[:, 1] - franka_limits[:, 0]
    ) * (limits[1] - limits[0]) + limits[0]
    return normalized


def _normalize_franka_joints_torch(
    batch_trajectory: torch.Tensor,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> torch.Tensor:
    """
    Normalizes joint angles to be within a specified range according to the Franka's
    joint limits. This is the torch version

    :param batch_trajectory torch.Tensor: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The new limits to map to
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype torch.Tensor: A tensor with the same dimensions as the input
    """
    assert isinstance(batch_trajectory, torch.Tensor)
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = torch.as_tensor(robot.JOINT_LIMITS).type_as(batch_trajectory)
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.size(0) == robot.DOF)
        or (batch_trajectory.ndim == 2 and batch_trajectory.size(1) == robot.DOF)
        or (batch_trajectory.ndim == 3 and batch_trajectory.size(2) == robot.DOF)
    )
    return (batch_trajectory - franka_limits[:, 0]) / (
        franka_limits[:, 1] - franka_limits[:, 0]
    ) * (limits[1] - limits[0]) + limits[0]


def normalize_franka_joints(
    batch_trajectory: Union[np.ndarray, torch.Tensor],
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> Union[np.ndarray, torch.Tensor]:
    """
    Normalizes joint angles to be within a specified range according to the Franka's
    joint limits. This is semantic sugar that dispatches to the correct implementation.

    :param batch_trajectory Union[np.ndarray, torch.Tensor]: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The new limits to map to
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype Union[np.ndarray, torch.Tensor]: A tensor or numpy array with the same dimensions
                                            and type as the input
    :raises NotImplementedError: Raises an error if another data type (e.g. a list) is passed in
    """
    if isinstance(batch_trajectory, torch.Tensor):
        return _normalize_franka_joints_torch(
            batch_trajectory, limits=limits, use_real_constraints=True
        )
    elif isinstance(batch_trajectory, np.ndarray):
        return _normalize_franka_joints_numpy(
            batch_trajectory, limits=limits, use_real_constraints=True
        )
    else:
        raise NotImplementedError("Only torch.Tensor and np.ndarray implemented")


def _unnormalize_franka_joints_numpy(
    batch_trajectory: np.ndarray,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> np.ndarray:
    """
    Unnormalizes joint angles from a specified range back into the Franka's joint limits.
    This is the numpy version and the inverse of `_normalize_franka_joints_numpy`.

    :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The current limits to map to the joint limits
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype np.ndarray: An array with the same dimensions as the input
    """
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = robot.JOINT_LIMITS
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF)
        or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF)
        or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF)
    )
    assert np.all(batch_trajectory >= limits[0])
    assert np.all(batch_trajectory <= limits[1])
    franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0]
    franka_lower_limit = franka_limits[:, 0]
    for _ in range(batch_trajectory.ndim - 1):
        franka_limit_range = franka_limit_range[np.newaxis, ...]
        franka_lower_limit = franka_lower_limit[np.newaxis, ...]
    unnormalized = (batch_trajectory - limits[0]) * franka_limit_range / (
        limits[1] - limits[0]
    ) + franka_lower_limit

    return unnormalized


def _unnormalize_franka_joints_torch(
    batch_trajectory: torch.Tensor,
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> torch.Tensor:
    """
    Unnormalizes joint angles from a specified range back into the Franka's joint limits.
    This is the torch version and the inverse of `_normalize_franka_joints_torch`.

    :param batch_trajectory torch.Tensor: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The current limits to map to the joint limits
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype torch.Tensor: A tensor with the same dimensions as the input
    """
    assert isinstance(batch_trajectory, torch.Tensor)
    robot = FrankaRealRobot if use_real_constraints else FrankaRobot
    franka_limits = torch.as_tensor(robot.JOINT_LIMITS).type_as(batch_trajectory)
    dof = franka_limits.size(0)
    assert (
        (batch_trajectory.ndim == 1 and batch_trajectory.size(0) == dof)
        or (batch_trajectory.ndim == 2 and batch_trajectory.size(1) == dof)
        or (batch_trajectory.ndim == 3 and batch_trajectory.size(2) == dof)
    )
    assert torch.all(batch_trajectory >= limits[0])
    assert torch.all(batch_trajectory <= limits[1])
    franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0]
    franka_lower_limit = franka_limits[:, 0]
    for _ in range(batch_trajectory.ndim - 1):
        franka_limit_range = franka_limit_range.unsqueeze(0)
        franka_lower_limit = franka_lower_limit.unsqueeze(0)
    return (batch_trajectory - limits[0]) * franka_limit_range / (
        limits[1] - limits[0]
    ) + franka_lower_limit


def unnormalize_franka_joints(
    batch_trajectory: Union[np.ndarray, torch.Tensor],
    limits: Tuple[float, float] = (-1, 1),
    use_real_constraints: bool = True,
) -> Union[np.ndarray, torch.Tensor]:
    """
    Unnormalizes joint angles from a specified range back into the Franka's joint limits.
    This is semantic sugar that dispatches to the correct implementation, the inverse of
    `normalize_franka_joints`.

    :param batch_trajectory Union[np.ndarray, torch.Tensor]: A batch of trajectories. Can have dims
                                        [7] if a single configuration
                                        [B, 7] if a batch of configurations
                                        [B, T, 7] if a batched time-series of configurations
    :param limits Tuple[float, float]: The current limits to map to the joint limits
    :param use_real_constraints bool: If true, use the empirically determined joint limits
                                      (this is unpublished--just found by monkeying around
                                      with the robot).
                                      If false, use the published joint limits from Franka
    :rtype Union[np.ndarray, torch.Tensor]: A tensor or numpy array with the same dimensions
                                            and type as the input
    :raises NotImplementedError: Raises an error if another data type (e.g. a list) is passed in
    """
    if isinstance(batch_trajectory, torch.Tensor):
        return _unnormalize_franka_joints_torch(
            batch_trajectory, limits=limits, use_real_constraints=use_real_constraints
        )
    elif isinstance(batch_trajectory, np.ndarray):
        return _unnormalize_franka_joints_numpy(
            batch_trajectory, limits=limits, use_real_constraints=use_real_constraints
        )
    else:
        raise NotImplementedError("Only torch.Tensor and np.ndarray implemented")

mpinets/utils.py
vishal-2000/EDMP