utils.py

#@markdown ### **Imports**
# diffusion policy import
from typing import Tuple, Sequence, Dict, Union, Optional, Callable
import numpy as np
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision.transforms import v2
import collections
import zarr
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
from diffusers.training_utils import EMAModel
from diffusers.optimization import get_scheduler
from tqdm.auto import tqdm
from PIL import Image
import pickle
# env import
import os


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def create_injected_noise(num_train_timesteps:int, beta_schedule='squaredcos_cap_v2'):
    noise_scheduler = DDPMScheduler(
        num_train_timesteps=num_train_timesteps,
        # the choise of beta schedule has big impact on performance
        # we found squared cosine works the best
        beta_schedule=beta_schedule,
        # clip output to [-1,1] to improve stability
        clip_sample=True,
        # our network predicts noise (instead of denoised action)
        prediction_type='epsilon'
    )
    return noise_scheduler

def split_batch_by_id(batch, unique_ids):
    split_batches = []

    for unique_id in unique_ids:
        indices = torch.where(batch['id'] == unique_id)[0]
        mini_batch = {
            'image': batch['image'][indices],
            'agent_pos': batch['agent_pos'][indices],
            'action': batch['action'][indices],
            'id': batch['id'][indices]
        }
        split_batches.append(mini_batch)

    return split_batches

def save(ema, nets, models_save_dir):
    if not os.path.exists(models_save_dir):
        os.makedirs(models_save_dir)
    torch.save(ema.state_dict(), os.path.join(models_save_dir, "ema_nets.pth"))
    for model_name, model in nets.items():
        model_path = os.path.join(models_save_dir, f"{model_name}.pth")
        torch.save(model.state_dict(), model_path)
        print(f"{model_name}.pth saved")

    print("All models have been saved successfully.")


#@markdown ### **Dataset**
#@markdown
#@markdown Defines `PushTImageDataset` and helper functions
#@markdown
#@markdown The dataset class
#@markdown - Load data ((image, agent_pos), action) from a zarr storage
#@markdown - Normalizes each dimension of agent_pos and action to [-1,1]
#@markdown - Returns
#@markdown  - All possible segments with length `pred_horizon`
#@markdown  - Pads the beginning and the end of each episode with repetition
#@markdown  - key `image`: shape (obs_hoirzon, 3, 96, 96)
#@markdown  - key `agent_pos`: shape (obs_hoirzon, 2)
#@markdown  - key `action`: shape (pred_horizon, 2)

def create_sample_indices(
        episode_ends:np.ndarray, sequence_length:int,
        pad_before: int=0, pad_after: int=0):
    indices = list()
    for i in range(len(episode_ends)):
        start_idx = 0
        if i > 0:
            start_idx = episode_ends[i-1]
        end_idx = episode_ends[i]
        episode_length = end_idx - start_idx

        min_start = -pad_before
        max_start = episode_length - sequence_length + pad_after

        # range stops one idx before end
        for idx in range(min_start, max_start+1):
            buffer_start_idx = max(idx, 0) + start_idx
            buffer_end_idx = min(idx+sequence_length, episode_length) + start_idx
            start_offset = buffer_start_idx - (idx+start_idx)
            end_offset = (idx+sequence_length+start_idx) - buffer_end_idx
            sample_start_idx = 0 + start_offset
            sample_end_idx = sequence_length - end_offset
            indices.append([
                buffer_start_idx, buffer_end_idx,
                sample_start_idx, sample_end_idx])
    indices = np.array(indices)
    return indices


def sample_sequence(train_data, sequence_length,
                    buffer_start_idx, buffer_end_idx,
                    sample_start_idx, sample_end_idx):
    result = dict()
    for key, input_arr in train_data.items():
        sample = input_arr[buffer_start_idx:buffer_end_idx]
        data = sample
        if (sample_start_idx > 0) or (sample_end_idx < sequence_length):
            data = np.zeros(
                shape=(sequence_length,) + input_arr.shape[1:],
                dtype=input_arr.dtype)
            if sample_start_idx > 0:
                data[:sample_start_idx] = sample[0]
            if sample_end_idx < sequence_length:
                data[sample_end_idx:] = sample[-1]
            data[sample_start_idx:sample_end_idx] = sample
        result[key] = data
    return result

# normalize data
def get_data_stats(data):
    data = data.reshape(-1,data.shape[-1])
    stats = {
        'min': np.min(data, axis=0),
        'max': np.max(data, axis=0)
    }
    return stats

def normalize_data(data, stats):
    # nomalize to [0,1]
    ndata = (data - stats['min']) / (stats['max'] - stats['min'])
    # normalize to [-1, 1]
    ndata = ndata * 2 - 1
    return ndata

def unnormalize_data(ndata, stats):
    ndata = (ndata + 1) / 2
    data = ndata * (stats['max'] - stats['min']) + stats['min']
    return data


import clip
from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path
from llava.eval.run_llava import my_eval_model
# Load the CLIP model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device)

# Encode a sentence into a feature vector
def encode_text(sentence):
    #sentence = truncate_sentence(sentence)
    text = clip.tokenize([sentence]).to(device)
    with torch.no_grad():
        text_features = clip_model.encode_text(text)
    return text_features


# dataset
class TrainDataset(torch.utils.data.Dataset):
    def __init__(self,
                 dataset_path: str,
                 pred_horizon: int,
                 obs_horizon: int,
                 action_horizon: int,
                 id:int,
                 num_demos: int,
                 resize_scale: int, 
                 pretrained=False,
                 stats = None):

        # read from zarr dataset
        dataset_root = np.load(dataset_path, allow_pickle=True).item()


        if 'outbin' in dataset_path:
            language_instruction = 'place blue block into small white box'
        else:
            language_instruction = 'remove red block from small white box'

        language_instruction_embedding = encode_text(language_instruction)
        language_instruction_embedding = language_instruction_embedding / language_instruction_embedding.norm(dim=-1, keepdim=True)
        language_instruction_embedding = language_instruction_embedding[0].cpu().numpy()


        end_frames = dataset_root['end_frames']
        all_actions = dataset_root['all_actions']
        all_agent_pose = dataset_root['all_agent_pose']
        
        all_pointcloud = dataset_root['all_pointcloud']
        relationship = dataset_root['all_relationship']
        
        num_max_demos = end_frames.shape[0]
        num_demos = min(num_max_demos, num_demos)
        num_max_frames = end_frames[num_demos-1]

        episode_ends = end_frames[:num_demos]

        # (N, D)
        train_data = {
            # first two dims of state vector are agent (i.e. gripper) locations
            'agent_pos': all_agent_pose[:num_max_frames],
            'action': all_actions[:num_max_frames]
        }

        # compute start and end of each state-action sequence
        # also handles padding
        indices = create_sample_indices(
            episode_ends=episode_ends,
            sequence_length=pred_horizon,
            pad_before=obs_horizon-1,
            pad_after=action_horizon-1)

        if stats == None:
            stats = dict()
            normalized_train_data = dict()
            for key, data in train_data.items():
                stats[key] = get_data_stats(data)
                normalized_train_data[key] = normalize_data(data, stats[key])
        else:
            # compute statistics and normalized data to [-1,1]
            normalized_train_data = dict()
            for key, data in train_data.items():
                normalized_train_data[key] = normalize_data(data, stats[key])

        normalized_train_data['point_cloud_0'] = np.array(all_pointcloud[0])
        normalized_train_data['point_cloud_1'] = np.array(all_pointcloud[1])
        normalized_train_data['point_cloud_2'] = np.array(all_pointcloud[2])
        
        normalized_train_data['relationship'] = np.array(relationship)

        normalized_train_data['instruction'] = np.tile(np.array(language_instruction_embedding), (normalized_train_data['relationship'].shape[0], 1))
        
        
        self.indices = indices
        self.stats = stats
        self.normalized_train_data = normalized_train_data
        self.pred_horizon = pred_horizon
        self.action_horizon = action_horizon
        self.obs_horizon = obs_horizon
        self.dataset_path = dataset_path
        self.id = id
        self.resize_scale = resize_scale
        self.pretrained = pretrained

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, idx):
        # get the start/end indices for this datapoint
        buffer_start_idx, buffer_end_idx, \
            sample_start_idx, sample_end_idx = self.indices[idx]

        # get nomralized data using these indices
        nsample = sample_sequence(
            train_data=self.normalized_train_data,
            sequence_length=self.pred_horizon,
            buffer_start_idx=buffer_start_idx,
            buffer_end_idx=buffer_end_idx,
            sample_start_idx=sample_start_idx,
            sample_end_idx=sample_end_idx
        )
        point_cloud_0 = nsample['point_cloud_0'][:self.obs_horizon,:]
        point_cloud_1 = nsample['point_cloud_1'][:self.obs_horizon,:]
        point_cloud_2 = nsample['point_cloud_2'][:self.obs_horizon,:]

        relationship = nsample['relationship'][:self.obs_horizon,:]
        instruction = nsample['instruction'][:self.obs_horizon,:]
    

        nsample['point_cloud_0'] = point_cloud_0
        nsample['point_cloud_1'] = point_cloud_1
        nsample['point_cloud_2'] = point_cloud_2

        
        nsample['relationship'] = relationship
        nsample['instruction'] = instruction

        nsample['agent_pos'] = nsample['agent_pos'][:self.obs_horizon,:]

        
        return nsample