492 lines
16 KiB
Python
492 lines
16 KiB
Python
import logging
|
|
|
|
import os
|
|
import glob
|
|
import torch
|
|
import pickle
|
|
import numpy as np
|
|
|
|
try:
|
|
import cv2 # not included in pyproject.toml
|
|
except:
|
|
print("Installing cv2")
|
|
os.system("pip install opencv-python")
|
|
from tqdm import tqdm
|
|
|
|
from agent.dataset.d3il_dataset.base_dataset import TrajectoryDataset
|
|
from agent.dataset.d3il_dataset import sim_framework_path
|
|
from .geo_transform import quat2euler
|
|
|
|
|
|
class Sorting_Dataset(TrajectoryDataset):
|
|
def __init__(
|
|
self,
|
|
data_directory: os.PathLike,
|
|
device="cpu",
|
|
obs_dim: int = 20,
|
|
action_dim: int = 2,
|
|
max_len_data: int = 256,
|
|
window_size: int = 1,
|
|
num_boxes: int = 2,
|
|
):
|
|
|
|
super().__init__(
|
|
data_directory=data_directory,
|
|
device=device,
|
|
obs_dim=obs_dim,
|
|
action_dim=action_dim,
|
|
max_len_data=max_len_data,
|
|
window_size=window_size,
|
|
)
|
|
|
|
logging.info("Loading Sorting Dataset")
|
|
|
|
inputs = []
|
|
actions = []
|
|
masks = []
|
|
|
|
# for root, dirs, files in os.walk(self.data_directory):
|
|
#
|
|
# for mode_dir in dirs:
|
|
|
|
# state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
|
|
# data_dir = os.path.join(sim_framework_path(data_directory), "local")
|
|
# data_dir = sim_framework_path(data_directory)
|
|
# state_files = glob.glob(data_dir + "/env*")
|
|
|
|
# random.seed(0)
|
|
|
|
# data_dir = sim_framework_path(data_directory)
|
|
# state_files = os.listdir(data_dir)
|
|
#
|
|
# random.shuffle(state_files)
|
|
#
|
|
# if data == "train":
|
|
# env_state_files = state_files[50:]
|
|
# elif data == "eval":
|
|
# env_state_files = state_files[:50]
|
|
# else:
|
|
# assert False, "wrong data type"
|
|
|
|
if num_boxes == 2:
|
|
data_dir = sim_framework_path("data/sorting/2_boxes/state")
|
|
elif num_boxes == 4:
|
|
data_dir = sim_framework_path("data/sorting/4_boxes/state")
|
|
elif num_boxes == 6:
|
|
data_dir = sim_framework_path("data/sorting/6_boxes/state")
|
|
else:
|
|
assert False, "check num boxes"
|
|
|
|
state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
|
|
|
|
for file in state_files:
|
|
with open(os.path.join(data_dir, file), "rb") as f:
|
|
env_state = pickle.load(f)
|
|
|
|
# lengths.append(len(env_state['robot']['des_c_pos']))
|
|
|
|
zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
|
|
zero_action = np.zeros(
|
|
(1, self.max_len_data, self.action_dim), dtype=np.float32
|
|
)
|
|
zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
|
|
|
|
# robot and box posistion
|
|
robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
|
|
robot_c_pos = env_state["robot"]["c_pos"][:, :2]
|
|
|
|
if num_boxes == 2:
|
|
red_box1_pos = env_state["red-box1"]["pos"][:, :2]
|
|
red_box1_quat = np.tan(
|
|
quat2euler(env_state["red-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
|
|
blue_box1_quat = np.tan(
|
|
quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
input_state = np.concatenate(
|
|
(
|
|
robot_des_pos,
|
|
robot_c_pos,
|
|
red_box1_pos,
|
|
red_box1_quat,
|
|
blue_box1_pos,
|
|
blue_box1_quat,
|
|
),
|
|
axis=-1,
|
|
)
|
|
|
|
elif num_boxes == 4:
|
|
|
|
red_box1_pos = env_state["red-box1"]["pos"][:, :2]
|
|
red_box1_quat = np.tan(
|
|
quat2euler(env_state["red-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
red_box2_pos = env_state["red-box2"]["pos"][:, :2]
|
|
red_box2_quat = np.tan(
|
|
quat2euler(env_state["red-box2"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
|
|
blue_box1_quat = np.tan(
|
|
quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box2_pos = env_state["blue-box2"]["pos"][:, :2]
|
|
blue_box2_quat = np.tan(
|
|
quat2euler(env_state["blue-box2"]["quat"])[:, -1:]
|
|
)
|
|
|
|
input_state = np.concatenate(
|
|
(
|
|
robot_des_pos,
|
|
robot_c_pos,
|
|
red_box1_pos,
|
|
red_box1_quat,
|
|
red_box2_pos,
|
|
red_box2_quat,
|
|
blue_box1_pos,
|
|
blue_box1_quat,
|
|
blue_box2_pos,
|
|
blue_box2_quat,
|
|
),
|
|
axis=-1,
|
|
)
|
|
|
|
elif num_boxes == 6:
|
|
|
|
red_box1_pos = env_state["red-box1"]["pos"][:, :2]
|
|
red_box1_quat = np.tan(
|
|
quat2euler(env_state["red-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
red_box2_pos = env_state["red-box2"]["pos"][:, :2]
|
|
red_box2_quat = np.tan(
|
|
quat2euler(env_state["red-box2"]["quat"])[:, -1:]
|
|
)
|
|
|
|
red_box3_pos = env_state["red-box3"]["pos"][:, :2]
|
|
red_box3_quat = np.tan(
|
|
quat2euler(env_state["red-box3"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box1_pos = env_state["blue-box1"]["pos"][:, :2]
|
|
blue_box1_quat = np.tan(
|
|
quat2euler(env_state["blue-box1"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box2_pos = env_state["blue-box2"]["pos"][:, :2]
|
|
blue_box2_quat = np.tan(
|
|
quat2euler(env_state["blue-box2"]["quat"])[:, -1:]
|
|
)
|
|
|
|
blue_box3_pos = env_state["blue-box3"]["pos"][:, :2]
|
|
blue_box3_quat = np.tan(
|
|
quat2euler(env_state["blue-box3"]["quat"])[:, -1:]
|
|
)
|
|
|
|
input_state = np.concatenate(
|
|
(
|
|
robot_des_pos,
|
|
robot_c_pos,
|
|
red_box1_pos,
|
|
red_box1_quat,
|
|
red_box2_pos,
|
|
red_box2_quat,
|
|
red_box3_pos,
|
|
red_box3_quat,
|
|
blue_box1_pos,
|
|
blue_box1_quat,
|
|
blue_box2_pos,
|
|
blue_box2_quat,
|
|
blue_box3_pos,
|
|
blue_box3_quat,
|
|
),
|
|
axis=-1,
|
|
)
|
|
|
|
else:
|
|
assert False, "check num boxes"
|
|
|
|
vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
|
|
valid_len = len(vel_state)
|
|
|
|
zero_obs[0, :valid_len, :] = input_state[:-1]
|
|
zero_action[0, :valid_len, :] = vel_state
|
|
zero_mask[0, :valid_len] = 1
|
|
|
|
inputs.append(zero_obs)
|
|
actions.append(zero_action)
|
|
masks.append(zero_mask)
|
|
|
|
# shape: B, T, n
|
|
self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
|
|
self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
|
|
self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
|
|
|
|
self.num_data = len(self.observations)
|
|
|
|
self.slices = self.get_slices()
|
|
|
|
def get_slices(self):
|
|
slices = []
|
|
|
|
min_seq_length = np.inf
|
|
for i in range(self.num_data):
|
|
T = self.get_seq_length(i)
|
|
min_seq_length = min(T, min_seq_length)
|
|
|
|
if T - self.window_size < 0:
|
|
print(
|
|
f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
|
|
)
|
|
else:
|
|
slices += [
|
|
(i, start, start + self.window_size)
|
|
for start in range(T - self.window_size + 1)
|
|
] # slice indices follow convention [start, end)
|
|
|
|
return slices
|
|
|
|
def get_seq_length(self, idx):
|
|
return int(self.masks[idx].sum().item())
|
|
|
|
def get_all_actions(self):
|
|
result = []
|
|
# mask out invalid actions
|
|
for i in range(len(self.masks)):
|
|
T = int(self.masks[i].sum().item())
|
|
result.append(self.actions[i, :T, :])
|
|
return torch.cat(result, dim=0)
|
|
|
|
def get_all_observations(self):
|
|
result = []
|
|
# mask out invalid observations
|
|
for i in range(len(self.masks)):
|
|
T = int(self.masks[i].sum().item())
|
|
result.append(self.observations[i, :T, :])
|
|
return torch.cat(result, dim=0)
|
|
|
|
def __len__(self):
|
|
return len(self.slices)
|
|
|
|
def __getitem__(self, idx):
|
|
|
|
i, start, end = self.slices[idx]
|
|
|
|
obs = self.observations[i, start:end]
|
|
act = self.actions[i, start:end]
|
|
mask = self.masks[i, start:end]
|
|
|
|
return obs, act, mask
|
|
|
|
|
|
class Sorting_Img_Dataset(TrajectoryDataset):
|
|
def __init__(
|
|
self,
|
|
data_directory: os.PathLike,
|
|
device="cpu",
|
|
obs_dim: int = 20,
|
|
action_dim: int = 2,
|
|
max_len_data: int = 256,
|
|
window_size: int = 1,
|
|
num_boxes: int = 2,
|
|
):
|
|
|
|
super().__init__(
|
|
data_directory=data_directory,
|
|
device=device,
|
|
obs_dim=obs_dim,
|
|
action_dim=action_dim,
|
|
max_len_data=max_len_data,
|
|
window_size=window_size,
|
|
)
|
|
|
|
logging.info("Loading Sorting Dataset")
|
|
|
|
inputs = []
|
|
actions = []
|
|
masks = []
|
|
|
|
# for root, dirs, files in os.walk(self.data_directory):
|
|
#
|
|
# for mode_dir in dirs:
|
|
|
|
# state_files = glob.glob(os.path.join(root, mode_dir) + "/env*")
|
|
# data_dir = os.path.join(sim_framework_path(data_directory), "local")
|
|
# data_dir = sim_framework_path(data_directory)
|
|
# state_files = glob.glob(data_dir + "/env*")
|
|
|
|
# random.seed(0)
|
|
#
|
|
# data_dir = sim_framework_path(data_directory)
|
|
# state_files = glob.glob(data_dir + '/state/*')
|
|
#
|
|
# random.shuffle(state_files)
|
|
#
|
|
# if data == "train":
|
|
# env_state_files = state_files[30:]
|
|
# elif data == "eval":
|
|
# env_state_files = state_files[:30]
|
|
# else:
|
|
# assert False, "wrong data type"
|
|
|
|
if num_boxes == 2:
|
|
data_dir = sim_framework_path("environments/dataset/data/sorting/2_boxes/")
|
|
elif num_boxes == 4:
|
|
data_dir = sim_framework_path("environments/dataset/data/sorting/4_boxes/")
|
|
elif num_boxes == 6:
|
|
data_dir = sim_framework_path("environments/dataset/data/sorting/6_boxes/")
|
|
else:
|
|
assert False, "check num boxes"
|
|
|
|
state_files = np.load(sim_framework_path(data_directory), allow_pickle=True)
|
|
|
|
bp_cam_imgs = []
|
|
inhand_cam_imgs = []
|
|
|
|
for file in tqdm(state_files[:100]):
|
|
with open(os.path.join(data_dir, "state", file), "rb") as f:
|
|
env_state = pickle.load(f)
|
|
|
|
# lengths.append(len(env_state['robot']['des_c_pos']))
|
|
zero_obs = np.zeros((1, self.max_len_data, self.obs_dim), dtype=np.float32)
|
|
zero_action = np.zeros(
|
|
(1, self.max_len_data, self.action_dim), dtype=np.float32
|
|
)
|
|
zero_mask = np.zeros((1, self.max_len_data), dtype=np.float32)
|
|
|
|
# robot and box posistion
|
|
robot_des_pos = env_state["robot"]["des_c_pos"][:, :2]
|
|
robot_c_pos = env_state["robot"]["c_pos"][:, :2]
|
|
|
|
file_name = os.path.basename(file).split(".")[0]
|
|
|
|
###############################################################
|
|
bp_images = []
|
|
bp_imgs = glob.glob(data_dir + "/images/bp-cam/" + file_name + "/*")
|
|
bp_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
|
|
|
|
for img in bp_imgs:
|
|
image = cv2.imread(img).astype(np.float32)
|
|
image = image.transpose((2, 0, 1)) / 255.0
|
|
|
|
image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
|
|
|
|
bp_images.append(image)
|
|
|
|
bp_images = torch.concatenate(bp_images, dim=0)
|
|
################################################################
|
|
inhand_imgs = glob.glob(data_dir + "/images/inhand-cam/" + file_name + "/*")
|
|
inhand_imgs.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
|
|
inhand_images = []
|
|
for img in inhand_imgs:
|
|
image = cv2.imread(img).astype(np.float32)
|
|
image = image.transpose((2, 0, 1)) / 255.0
|
|
|
|
image = torch.from_numpy(image).to(self.device).float().unsqueeze(0)
|
|
|
|
inhand_images.append(image)
|
|
inhand_images = torch.concatenate(inhand_images, dim=0)
|
|
##################################################################
|
|
# input_state = np.concatenate((robot_des_pos, robot_c_pos), axis=-1)
|
|
|
|
vel_state = robot_des_pos[1:] - robot_des_pos[:-1]
|
|
|
|
valid_len = len(vel_state)
|
|
|
|
zero_obs[0, :valid_len, :] = robot_des_pos[:-1]
|
|
zero_action[0, :valid_len, :] = vel_state
|
|
zero_mask[0, :valid_len] = 1
|
|
|
|
bp_cam_imgs.append(bp_images)
|
|
inhand_cam_imgs.append(inhand_images)
|
|
|
|
inputs.append(zero_obs)
|
|
actions.append(zero_action)
|
|
masks.append(zero_mask)
|
|
|
|
self.bp_cam_imgs = bp_cam_imgs
|
|
self.inhand_cam_imgs = inhand_cam_imgs
|
|
|
|
# shape: B, T, n
|
|
self.observations = torch.from_numpy(np.concatenate(inputs)).to(device).float()
|
|
self.actions = torch.from_numpy(np.concatenate(actions)).to(device).float()
|
|
self.masks = torch.from_numpy(np.concatenate(masks)).to(device).float()
|
|
|
|
self.num_data = len(self.actions)
|
|
|
|
self.slices = self.get_slices()
|
|
|
|
def get_slices(self):
|
|
slices = []
|
|
|
|
min_seq_length = np.inf
|
|
for i in range(self.num_data):
|
|
T = self.get_seq_length(i)
|
|
min_seq_length = min(T, min_seq_length)
|
|
|
|
if T - self.window_size < 0:
|
|
print(
|
|
f"Ignored short sequence #{i}: len={T}, window={self.window_size}"
|
|
)
|
|
else:
|
|
slices += [
|
|
(i, start, start + self.window_size)
|
|
for start in range(T - self.window_size + 1)
|
|
] # slice indices follow convention [start, end)
|
|
|
|
return slices
|
|
|
|
def get_seq_length(self, idx):
|
|
return int(self.masks[idx].sum().item())
|
|
|
|
def get_all_actions(self):
|
|
result = []
|
|
# mask out invalid actions
|
|
for i in range(len(self.masks)):
|
|
T = int(self.masks[i].sum().item())
|
|
result.append(self.actions[i, :T, :])
|
|
return torch.cat(result, dim=0)
|
|
|
|
def get_all_observations(self):
|
|
result = []
|
|
# mask out invalid observations
|
|
for i in range(len(self.masks)):
|
|
T = int(self.masks[i].sum().item())
|
|
result.append(self.observations[i, :T, :])
|
|
return torch.cat(result, dim=0)
|
|
|
|
def __len__(self):
|
|
return len(self.slices)
|
|
|
|
def __getitem__(self, idx):
|
|
|
|
i, start, end = self.slices[idx]
|
|
|
|
obs = self.observations[i, start:end]
|
|
act = self.actions[i, start:end]
|
|
mask = self.masks[i, start:end]
|
|
|
|
bp_imgs = self.bp_cam_imgs[i][start:end]
|
|
inhand_imgs = self.inhand_cam_imgs[i][start:end]
|
|
|
|
# bp_imgs = np.zeros((self.window_size, 3, 96, 96), dtype=np.float32)
|
|
# inhand_imgs = np.zeros((self.window_size, 3, 96, 96), dtype=np.float32)
|
|
#
|
|
# for num_frame, img_file in enumerate(bp_img_files):
|
|
# image = cv2.imread(img_file).astype(np.float32)
|
|
# bp_imgs[num_frame] = image.transpose((2, 0, 1)) / 255.
|
|
#
|
|
# for num_frame, img_file in enumerate(inhand_img_files):
|
|
# image = cv2.imread(img_file).astype(np.float32)
|
|
# inhand_imgs[num_frame] = image.transpose((2, 0, 1)) / 255.
|
|
#
|
|
# bp_imgs = torch.from_numpy(bp_imgs).to(self.device).float()
|
|
# inhand_imgs = torch.from_numpy(inhand_imgs).to(self.device).float()
|
|
|
|
return bp_imgs, inhand_imgs, obs, act, mask
|