check_observation_poses.py

import numpy as np
import torch
import trimesh

import pickpp

import contrib


np.random.seed(0)

models = pickpp.datasets.YCBVideoModelDataset()

env = contrib.RlightPickAndPlaceEnv(version="v1", headless=False)
obs = env.reset()

cads = torch.as_tensor(obs["cads"][None])
objects = torch.as_tensor(obs["objects"][None])
robot = torch.as_tensor(obs["robot"][None])

B, O = objects.shape[:2]

world = False

translation = objects[:, :, 1:4].reshape(-1, 3)
quaternion = objects[:, :, 4:8].reshape(-1, 4)
T_cad2robot = pickpp.functions.transformation_matrix(
    quaternion, translation
)
T_ee2robot = pickpp.functions.transformation_matrix(
    quaternion=robot[:, 3:], translation=robot[:, :3],
)
if world:
    pass
else:
    T_ee2robot = (
        T_ee2robot[:, None]
        .repeat_interleave(O, dim=1)
        .reshape(-1, 4, 4)
    )
    T_robot2ee = torch.inverse(T_ee2robot)
    T_cad2ee = T_robot2ee @ T_cad2robot
    # translation = pickpp.functions.translation_from_matrix(T_cad2ee)
    # quaternion = pickpp.functions.quaternion_from_matrix(T_cad2ee)

scene = trimesh.Scene()

# cad
for o in range(O):
    cad_id = int(cads[0, o].item())
    if cad_id == 0:
        continue
    cad_file = models.get_cad_file_from_id(cad_id=cad_id)
    cad = trimesh.load(cad_file)
    cad.visual = cad.visual.to_color()
    axis = trimesh.creation.axis(origin_size=0.01)
    if world:
        cad.apply_transform(T_cad2robot[o].numpy())
        scene.add_geometry(cad)
        axis.apply_transform(T_cad2robot[o].numpy())
        scene.add_geometry(axis)
    else:
        cad.apply_transform(T_cad2ee[o].numpy())
        scene.add_geometry(cad)
        axis.apply_transform(T_cad2ee[o].numpy())
        scene.add_geometry(axis)

# base
axis = trimesh.creation.axis(origin_size=0.01, origin_color=(0, 0, 0))
scene.add_geometry(axis)

axis = trimesh.creation.axis(origin_size=0.01, origin_color=(0.5, 0.5, 0.5))
if world:
    axis.apply_transform(T_ee2robot[0].numpy())
    scene.add_geometry(axis)
else:
    axis.apply_transform(T_robot2ee[0].numpy())
    scene.add_geometry(axis)

scene.show()