def chart_current_v_next(demo_path):
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
demo_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=0,
to_demo=1)
# Load all of the poses into a list, reuse some load-data functions here
demo_trajectories = [
get_pose_hist(d, "l_wrist_roll_link") for d in demo_paths
]
demo_trajectories = [
np.array([quat_pose_to_rpy(p) for p in dt]).transpose(1, 0)
for dt in demo_trajectories
]
pose_dim_names = ["x", "y", "z", "r-x", "r-y", "r-z"]
for i in range(6):
plt.subplot(2, 3, i + 1)
plt.xlabel("current {}".format(pose_dim_names[i]))
plt.ylabel("next {}".format(pose_dim_names[i]))
for d in demo_trajectories:
plt.scatter(d[i][0:-5], d[i][5:], alpha=0.05, s=10, color="blue")
plt.show()
def chart_all_poses_3d(demo_path):
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
demo_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=0,
to_demo=5)
# Load all of the poses into a list, reuse some load-data functions here
demo_trajectories = [
get_pose_hist(d, "l_wrist_roll_link") for d in demo_paths
]
demo_trajectories = [
np.array([quat_pose_to_rpy(p) for p in dt]).transpose(1, 0)
for dt in demo_trajectories
]
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("t")
for d in demo_trajectories:
time = np.linspace(0, 1, len(d[2]))
ax.scatter(d[0], d[1], time)
plt.show()
def chart_all_poses(demo_path):
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
demo_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=0)
# Load all of the poses into a list, reuse some load-data functions here
demo_trajectories = [
get_pose_hist(d, "l_wrist_roll_link") for d in demo_paths
]
# demo_trajectories = [np.array([quat_pose_to_rpy(p) for p in dt]).transpose(1,0) for dt in demo_trajectories]
smoothed_trajectories = []
for dt in demo_trajectories:
smoothed_dt = [dt[0]]
for i in range(1, len(dt)):
# print(i)
current_pos, current_quat = dt[i][0:3], dt[i][3:]
prev_quat = smoothed_dt[i - 1][3:]
orig_diff = np.abs(current_quat - prev_quat).sum()
flipped_diff = np.abs(-current_quat - prev_quat).sum()
if flipped_diff <= orig_diff:
print("Orig diff {}, Flipped diff{} at {} ".format(
orig_diff, flipped_diff, i))
print("Original quat {}, Flipped {} ".format(
current_quat, -current_quat))
smoothed_dt.append(np.concatenate(
(current_pos, -current_quat)))
print(len(smoothed_dt))
else:
smoothed_dt.append(np.concatenate((current_pos, current_quat)))
smoothed_trajectories.append(np.array(smoothed_dt))
# smoothed_trajectories = demo_trajectories
# demo_trajectories = [dt.transpose(1,0) for dt in demo_trajectories]
demo_trajectories = [dt.transpose(1, 0) for dt in smoothed_trajectories]
# pose_dim_names = ["x", "y", "z", "r-x", "r-y", "r-z"]
pose_dim_names = ["x", "y", "z", "q-x", "q-y", "q-z", "q-w"]
for i in range(len(pose_dim_names)):
plt.subplot(3, 3, i + 1)
plt.xlabel("time")
plt.ylabel((pose_dim_names[i]))
for d in demo_trajectories:
plt.plot(d[i], alpha=0.15, color="green")
plt.show()
def chart_mdn_means(model_path, demo_path, demo_num):
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
_, demo_loader = load_pose_state_demos(
image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=demo_num,
to_demo=demo_num + 1), exp_config["batch_size"],
"l_wrist_roll_link", "r_wrist_roll_link", False, torch.device("cuda"))
model = PosePlusStateNet(100, 2)
model.load_state_dict(
torch.load(model_path, map_location=torch.device("cuda")))
model.to(torch.device("cuda"))
model.eval()
mu_all = []
std_all = []
pis_all = []
targets_all = []
with torch.no_grad():
for ins in demo_loader:
current_pose, goal_pose, targets = ins
pis, stds, mus = model(current_pose, goal_pose)
print("mu {} std {} pis {}".format(mus.shape, stds.shape,
pis.shape))
mu_all.append(mus)
std_all.append(stds)
pis_all.append(pis)
targets_all.append(targets)
mu_all = torch.cat(mu_all)
std_all = torch.cat(std_all)
pis_all = torch.cat(pis_all)
targets_all = torch.cat(targets_all)
# Im expecting: mu: N x 5 X 7, std: N x 5, pis: N * 5, targets_all: N * 7
# print("Shapes: mu: {}, std: {}, pis {}, targets{}".format(mu_all.shape, std_all.shape, pis_all.shape, targets_all.shape))
# I want : mu: 7 X 5 X N, std: 5 X N, pis: 5 X N, targets_all: 7 * N
mu_all = mu_all.permute(2, 1, 0).cpu().numpy()
std_all = std_all.permute(2, 1, 0).cpu().numpy()
pis_all = pis_all.permute(1, 0).cpu().numpy()
targets_all = targets_all.permute(1, 0).cpu().numpy()
print("Shapes: mu: {}, std: {}, pis {}, targets{}".format(
mu_all.shape, std_all.shape, pis_all.shape, targets_all.shape))
fig = plt.figure()
print("STD example", std_all[0][0])
n_joints, n_components, n_samples = mu_all.shape[0], mu_all.shape[
1], mu_all.shape[2]
print("joints {} components {} samples {}".format(n_joints, n_components,
n_samples))
for joint_id in range(n_joints):
plt.subplot(ceil(n_joints / 3.0), 4, joint_id + 1)
for pi_id in range(n_components):
plt.plot(mu_all[joint_id, pi_id], label="pi-{}".format(pi_id))
plt.fill_between(
range(n_samples),
mu_all[joint_id, pi_id] - std_all[joint_id, pi_id],
mu_all[joint_id, pi_id] + std_all[joint_id, pi_id],
alpha=0.1)
weighted_mus = mu_all[joint_id] * pis_all
averaged_mus = weighted_mus.sum(0)
print("muall: {}, pis_all: {} Weighted: {}, Averaged: {}".format(
mu_all.shape, pis_all.shape, weighted_mus.shape,
averaged_mus.shape))
plt.plot(averaged_mus, label="Averaged")
plt.plot(targets_all[joint_id], label="Targets")
plt.legend()
plt.title(["x", "y", "z", "r-x", "r-y", "r-z"][joint_id])
plt.xlabel("t")
plt.ylabel("Normed Radians")
plt.subplot(ceil(n_joints / 3.0), 4, n_joints + 1)
for pi_id in range(n_components):
plt.plot(pis_all[pi_id], label="pi_{}".format(pi_id))
plt.legend()
plt.xlabel("t")
plt.ylabel("Component Weight")
plt.title("Pis")
plt.show()
def chart_pred_pose(model_path, demo_path, demo_num):
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
train_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=demo_num,
to_demo=demo_num + 1)
train_set = ImagePoseFuturePose(train_paths, "l_wrist_roll_link",
get_trans(im_params, distorted=True), 10)
demo_loader = DeviceDataLoader(
DataLoader(train_set, exp_config["batch_size"], shuffle=False),
torch.device("cuda"))
model = ImageOnlyNet(im_params["resize_height"], im_params["resize_width"],
6)
model.load_state_dict(
torch.load(model_path, map_location=torch.device("cuda")))
model.to(torch.device("cuda"))
model.eval()
next_pred_all = []
# current_pred_all = []
current_targets_all = []
next_targets_all = []
with torch.no_grad():
for ins in demo_loader:
rgb_ins, current_ins, target_ins = ins
next_pred = model(rgb_ins)
next_pred_all.append(next_pred)
# current_pred_all.append(current_pred)
current_targets_all.append(current_ins)
next_targets_all.append(target_ins)
next_pred_all = torch.cat(next_pred_all)
# current_pred_all = torch.cat(current_pred_all)
current_targets_all = torch.cat(current_targets_all)
next_targets_all = torch.cat(next_targets_all)
# Shapes: N X 7
next_pred_all = next_pred_all.permute(1, 0).cpu().numpy()
# current_pred_all = current_pred_all.permute(1,0).cpu().numpy()
current_targets_all = current_targets_all.permute(1, 0).cpu().numpy()
next_targets_all = next_targets_all.permute(1, 0).cpu().numpy()
n_pose_dims, n_samples = next_pred_all.shape[0], next_pred_all.shape[1]
print(n_pose_dims)
dim_names = ["p-x", "p-y", "p-z", "roll", "pitch", "yaw"]
for dim_id in range(n_pose_dims):
plt.subplot(ceil(n_pose_dims / 3.0), 3, dim_id + 1)
plt.plot(next_pred_all[dim_id], label="Predicted Pose")
# plt.plot(current_pred_all[dim_id], label="Aux Prediction")
plt.plot(current_targets_all[dim_id], label="Current Pose")
plt.plot(next_targets_all[dim_id], label="Target Pose")
plt.legend()
plt.xlabel("t")
plt.ylabel(dim_names[dim_id])
plt.show()
optimizer.zero_grad()
full_loss.backward()
optimizer.step()
return loss, constraint_loss, constraint_acc
if __name__ == "__main__":
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
model = ImagePlusPoseNet(
(im_params["resize_height"], im_params["resize_width"]), 100)
model.to(torch.device("cuda"))
train_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=0,
to_demo=75)
train_set = ImagePoseFuturePose(train_paths,
"l_wrist_roll_link",
get_trans(im_params, distorted=False),
skip_count=5)
train_loader = DeviceDataLoader(
DataLoader(train_set, exp_config["batch_size"], shuffle=True),
torch.device("cuda"))
# for i in range(0, len(train_set), 40):
# show_torched_im(train_set[i][0])
val_paths = image_demo_paths(exp_config["demo_folder"],
im_params["file_glob"],
from_demo=75)
aux_loss = F.l1_loss(aux_pred, target_aux)
full_loss = 1e-2 * l2_loss + 1.0 * l1_loss + 5e-3 * angle_loss + 1.0 * aux_loss
if math.isnan(full_loss):
print("Nan Encountered")
return full_loss, l2_loss, l1_loss, angle_loss, aux_loss
return lf
if __name__ == "__main__":
# Test out data loading accuracy
exp_config = load_json("config/experiment_config.json")
im_params = exp_config["image_config"]
train_set, train_loader = load_rgbd_demos(
image_demo_paths(exp_config["demo_folder"], im_params["file_glob"], from_demo=0, to_demo=90)
exp_config["batch_size"],
"l_wrist_roll_link",
get_trans(im_params, distorted=True),
get_grayscale_trans(im_params),
True,
torch.device("cuda"))
val_set, validation_loader = load_rgbd_demos(
image_demo_paths(exp_config["demo_folder"], im_params["file_glob"], from_demo=90)
exp_config["batch_size"],
"l_wrist_roll_link",
get_trans(im_params, distorted=True),
get_grayscale_trans(im_params),
False,
torch.device("cuda"))
