Hugging Face's logo

lsnu
/
VLAdaptorBench

Robotics
rlbench
benchmarking
label-validation
Model card Files
xet
 Community
VLAdaptorBench
File size: 4,525 Bytes
30747b3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
#!/usr/bin/env python3

from typing import List

import os

import numpy as np
import open3d as o3d

from rlbench import CameraConfig
from rlbench import ObservationConfig

from rlbench.utils import get_stored_demos
from pyrep.const import RenderMode

import rich_click as click
from click_prompt import filepath_option


box_length = 0.01
box_dim = dict(zip(["width", "height", "depth"], [box_length] * 3))
box_offset = np.identity(4)
box_offset[:3, 3] = - box_length / 2

camera_names = ["front", "overhead", "over_shoulder_left", "over_shoulder_right", "wrist_left", "wrist_right"]


@click.command()
@filepath_option("--task-folder", default="/tmp/rlbench_data/bimanual_push_box")
@click.option("--show-visualization/--hide-visualization", is_flag=True, default=True)
@click.option("--add-trajectory/--hide-trajectory", is_flag=True, default=True)
@click.option("--episode-number", "-e", help="Which episode to select from the dataset", default=0, type=int)
def cli(task_folder, show_visualization, add_trajectory, episode_number):

    task_folder = os.path.expanduser(task_folder)
    dataset_root = os.path.dirname(task_folder)
    task_name = os.path.basename(task_folder)

    obs_config = create_obs_config(camera_names, [128, 128])
    episodes = get_stored_demos(1, False, dataset_root, -1, task_name, obs_config, from_episode_number=episode_number)

    visualization_data = []


    # visualize trajectories
    if add_trajectory:
        for obs in episodes[episode_number]:
            right_action_box = o3d.geometry.TriangleMesh.create_box(**box_dim)
            left_action_box = o3d.geometry.TriangleMesh.create_box(**box_dim)

            right_action_box.paint_uniform_color([0, 1, 0]) 
            left_action_box.paint_uniform_color([1, 0, 0]) 

            right_action_box.transform(obs.right.gripper_matrix.dot(box_offset))
            left_action_box.transform(obs.left.gripper_matrix.dot(box_offset))

            visualization_data.append(right_action_box)
            visualization_data.append(left_action_box)



    # point cloud from first observation of the first episode
    obs = episodes[episode_number][0]
    for camera_name in camera_names:

        xyz = obs.perception_data[f"{camera_name}_point_cloud"].reshape(-1, 3)
        rgb = obs.perception_data[f"{camera_name}_rgb"].reshape(-1, 3) / 255.0
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(xyz)
        pcd.colors = o3d.utility.Vector3dVector(rgb)
        visualization_data.append(pcd)

    if show_visualization:

        o3d.visualization.draw_geometries(visualization_data)
    else:

        vis = o3d.visualization.Visualizer()
        vis.create_window(visible=False)
        for g in visualization_data:
            vis.add_geometry(g)
            vis.update_geometry(g)


        vis.poll_events()
        vis.update_renderer()

        view_control = vis.get_view_control()

        # copied from the viewer
        view_pose = {"front" : [ 0.012339452449784705, -0.90282996236235535, 0.42982065675584702 ],
                    "lookat" : [ -0.12694871669500873, 0.11746709358396157, 0.88075116287669197 ],
                    "up" : [ 0.025976937544841906, 0.42999774206857738, 0.90245617097547559 ],
                    "zoom" : 0.33999999999999964 }

        view_control.set_front(view_pose["front"])
        view_control.set_up(view_pose["up"])
        view_control.set_lookat(view_pose["lookat"])
        view_control.set_zoom(view_pose["zoom"])


        vis.capture_screen_image(f"/tmp/rlbench_{task_name}_episode{episode_number}.png", do_render=True)
        vis.destroy_window()


def create_obs_config(
    camera_names: List[str],
    camera_resolution: List[int],
):
    unused_cams = CameraConfig()
    unused_cams.set_all(False)
    used_cams = CameraConfig(
        rgb=True,
        point_cloud=True,
        mask=False,
        depth=False,
        image_size=camera_resolution,
        render_mode=RenderMode.OPENGL3,
    )

    camera_configs = {camera_name: used_cams for camera_name in camera_names}

    obs_config = ObservationConfig(
        camera_configs=camera_configs,
        joint_forces=False,
        joint_positions=True,
        joint_velocities=True,
        task_low_dim_state=False,
        gripper_touch_forces=False,
        gripper_pose=True,
        gripper_open=True,
        gripper_matrix=True,
        gripper_joint_positions=True,
        robot_name="bimanual"
    )
    return obs_config


if __name__ == "__main__":
    cli()