| |
|
|
| 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 = [] |
|
|
|
|
| |
| 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) |
|
|
|
|
|
|
| |
| 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() |
|
|
| |
| 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() |
|
|
|
|
