{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 可视化你的数据\n", "\n", "\n", "\n", "基于重建后的仿真场景可视化你的动作。\n", "\n", "主窗口会回放动作。\n", "\n", "右上和右下叠加图像来自数据集。\n" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "WARNING:root:'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "99d45a5ab77a492ab53ca0547ba9332d", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Generating train split: 0 examples [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata\n", "import numpy as np\n", "from lerobot.common.datasets.utils import write_json, serialize_dict\n", "\n", "dataset = LeRobotDataset('datawhale_eai_pnp', root='./demo_data') # if youu want to use the example data provided, root = './demo_data_example' instead!" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 加载数据集\n" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "class EpisodeSampler(torch.utils.data.Sampler):\n", " \"\"\"\n", " Sampler for a single episode\n", " \"\"\"\n", " def __init__(self, dataset: LeRobotDataset, episode_index: int):\n", " from_idx = dataset.episode_data_index[\"from\"][episode_index].item()\n", " to_idx = dataset.episode_data_index[\"to\"][episode_index].item()\n", " self.frame_ids = range(from_idx, to_idx)\n", "\n", " def __iter__(self):\n", " return iter(self.frame_ids)\n", "\n", " def __len__(self) -> int:\n", " return len(self.frame_ids)" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [], "source": [ "# Select an episode index that you want to visualize\n", "episode_index = 0\n", "\n", "episode_sampler = EpisodeSampler(dataset, episode_index)\n", "dataloader = torch.utils.data.DataLoader(\n", " dataset,\n", " num_workers=1,\n", " batch_size=1,\n", " sampler=episode_sampler,\n", ")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 在仿真中可视化数据集\n" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [], "source": [ "import env_config" ] }, { "cell_type": "code", "execution_count": 27, "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "-----------------------------------------------------------------------------\n", "name:[Tabletop] dt:[0.002] HZ:[500]\n", " n_qpos:[26] n_qvel:[24] n_qacc:[24] n_ctrl:[7]\n", " integrator:[RK4]\n", "\n", "n_body:[23]\n", " [0/23] [world] mass:[0.00]kg\n", " [1/23] [front_object_table] mass:[1.00]kg\n", " [2/23] [camera] mass:[0.00]kg\n", " [3/23] [camera2] mass:[0.00]kg\n", " [4/23] [camera3] mass:[0.00]kg\n", " [5/23] [nova5_mount] mass:[1.67]kg\n", " [6/23] [Link1] mass:[1.08]kg\n", " [7/23] [Link2] mass:[2.55]kg\n", " [8/23] [Link3] mass:[1.41]kg\n", " [9/23] [Link4] mass:[0.44]kg\n", " [10/23] [Link5] mass:[0.57]kg\n", " [11/23] [Link6] mass:[0.44]kg\n", " [12/23] [camera_wrist] mass:[0.00]kg\n", " [13/23] [left_outer_knuckle] mass:[0.03]kg\n", " [14/23] [left_inner_knuckle] mass:[0.01]kg\n", " [15/23] [left_inner_finger] mass:[0.01]kg\n", " [16/23] [right_inner_knuckle] mass:[0.01]kg\n", " [17/23] [right_inner_finger] mass:[0.01]kg\n", " [18/23] [right_outer_knuckle] mass:[0.03]kg\n", " [19/23] [body_obj_mug_5] mass:[0.00]kg\n", " [20/23] [object_mug_5] mass:[0.08]kg\n", " [21/23] [body_obj_plate_11] mass:[0.00]kg\n", " [22/23] [object_plate_11] mass:[0.10]kg\n", "body_total_mass:[9.45]kg\n", "\n", "n_geom:[84]\n", "geom_names:['floor', 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]\n", "\n", "n_mesh:[78]\n", "mesh_names:['base_link', 'Link1', 'Link2', 'Link3', 'Link4', 'Link5', 'Link6', 'robotiq_85_base_link_coarse', 'outer_knuckle_coarse', 'outer_finger_coarse', 'inner_knuckle_coarse', 'inner_finger_coarse', 'mug_5_normalized_0_vis', 'mug_5_normalized_collision_22._coll', 'mug_5_normalized_collision_23._coll', 'mug_5_normalized_collision_21._coll', 'mug_5_normalized_collision_20._coll', 'mug_5_normalized_collision_24._coll', 'mug_5_normalized_collision_30._coll', 'mug_5_normalized_collision_18._coll', 'mug_5_normalized_collision_19._coll', 'mug_5_normalized_collision_31._coll', 'mug_5_normalized_collision_25._coll', 'mug_5_normalized_collision_27._coll', 'mug_5_normalized_collision_26._coll', 'mug_5_normalized_collision_9._coll', 'mug_5_normalized_collision_8._coll', 'mug_5_normalized_collision_6._coll', 'mug_5_normalized_collision_7._coll', 'mug_5_normalized_collision_5._coll', 'mug_5_normalized_collision_4._coll', 'mug_5_normalized_collision_0._coll', 'mug_5_normalized_collision_1._coll', 'mug_5_normalized_collision_3._coll', 'mug_5_normalized_collision_2._coll', 'mug_5_normalized_collision_17._coll', 'mug_5_normalized_collision_16._coll', 'mug_5_normalized_collision_28._coll', 'mug_5_normalized_collision_14._coll', 'mug_5_normalized_collision_15._coll', 'mug_5_normalized_collision_29._coll', 'mug_5_normalized_collision_11._coll', 'mug_5_normalized_collision_10._coll', 'mug_5_normalized_collision_12._coll', 'mug_5_normalized_collision_13._coll', 'plate_11_normalized_0_vis', 'plate_11_normalized_collision_22._coll', 'plate_11_normalized_collision_23._coll', 'plate_11_normalized_collision_21._coll', 'plate_11_normalized_collision_20._coll', 'plate_11_normalized_collision_24._coll', 'plate_11_normalized_collision_30._coll', 'plate_11_normalized_collision_18._coll', 'plate_11_normalized_collision_19._coll', 'plate_11_normalized_collision_31._coll', 'plate_11_normalized_collision_25._coll', 'plate_11_normalized_collision_27._coll', 'plate_11_normalized_collision_26._coll', 'plate_11_normalized_collision_9._coll', 'plate_11_normalized_collision_8._coll', 'plate_11_normalized_collision_6._coll', 'plate_11_normalized_collision_7._coll', 'plate_11_normalized_collision_5._coll', 'plate_11_normalized_collision_4._coll', 'plate_11_normalized_collision_0._coll', 'plate_11_normalized_collision_1._coll', 'plate_11_normalized_collision_3._coll', 'plate_11_normalized_collision_2._coll', 'plate_11_normalized_collision_17._coll', 'plate_11_normalized_collision_16._coll', 'plate_11_normalized_collision_28._coll', 'plate_11_normalized_collision_14._coll', 'plate_11_normalized_collision_15._coll', 'plate_11_normalized_collision_29._coll', 'plate_11_normalized_collision_11._coll', 'plate_11_normalized_collision_10._coll', 'plate_11_normalized_collision_12._coll', 'plate_11_normalized_collision_13._coll']\n", "\n", "n_joint:[14]\n", " [0/14] [joint1] axis:[0. 0. 1.]\n", " [1/14] [joint2] axis:[0. 0. 1.]\n", " [2/14] [joint3] axis:[0. 0. 1.]\n", " [3/14] [joint4] axis:[0. 0. 1.]\n", " [4/14] [joint5] axis:[0. 0. 1.]\n", " [5/14] [joint6] axis:[0. 0. 1.]\n", " [6/14] [finger_joint] axis:[ 0. -1. 0.]\n", " [7/14] [left_inner_knuckle_joint] axis:[ 0. -1. 0.]\n", " [8/14] [left_inner_finger_joint] axis:[ 0. -1. 0.]\n", " [9/14] [right_inner_knuckle_joint] axis:[0. 1. 0.]\n", " [10/14] [right_inner_finger_joint] axis:[0. 1. 0.]\n", " [11/14] [right_outer_knuckle_joint] axis:[0. 1. 0.]\n", " [12/14] [None] axis:[0. 0. 1.]\n", " [13/14] [None] axis:[0. 0. 1.]\n", "\n", "n_dof:[24] (=number of rows of Jacobian)\n", " [0/24] [None] attached joint:[joint1] body:[Link1]\n", " [1/24] [None] attached joint:[joint2] body:[Link2]\n", " [2/24] [None] attached joint:[joint3] body:[Link3]\n", " [3/24] [None] attached joint:[joint4] body:[Link4]\n", " [4/24] [None] attached joint:[joint5] body:[Link5]\n", " [5/24] [None] attached joint:[joint6] body:[Link6]\n", " [6/24] [None] attached joint:[finger_joint] body:[left_outer_knuckle]\n", " [7/24] [None] attached joint:[left_inner_knuckle_joint] body:[left_inner_knuckle]\n", " [8/24] [None] attached joint:[left_inner_finger_joint] body:[left_inner_finger]\n", " [9/24] [None] attached joint:[right_inner_knuckle_joint] body:[right_inner_knuckle]\n", " [10/24] [None] attached joint:[right_inner_finger_joint] body:[right_inner_finger]\n", " [11/24] [None] attached joint:[right_outer_knuckle_joint] body:[right_outer_knuckle]\n", " [12/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [13/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [14/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [15/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [16/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [17/24] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [18/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [19/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [20/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [21/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [22/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [23/24] [None] attached joint:[None] body:[body_obj_plate_11]\n", "\n", "Free joint information. n_free_joint:[2]\n", " [0/2] [None] body_name_attached:[body_obj_mug_5]\n", " [1/2] [None] body_name_attached:[body_obj_plate_11]\n", "\n", "Revolute joint information. n_rev_joint:[12]\n", " [0/12] [joint1] range:[-6.280]~[6.280]\n", " [1/12] [joint2] range:[-3.140]~[3.140]\n", " [2/12] [joint3] range:[-2.790]~[2.790]\n", " [3/12] [joint4] range:[-6.280]~[6.280]\n", " [4/12] [joint5] range:[-6.280]~[6.280]\n", " [5/12] [joint6] range:[-6.280]~[6.280]\n", " [6/12] [finger_joint] range:[0.000]~[0.725]\n", " [7/12] [left_inner_knuckle_joint] range:[0.000]~[0.876]\n", " [8/12] [left_inner_finger_joint] range:[-0.876]~[0.000]\n", " [9/12] [right_inner_knuckle_joint] range:[-0.876]~[0.000]\n", " [10/12] [right_inner_finger_joint] range:[0.000]~[0.876]\n", " [11/12] [right_outer_knuckle_joint] range:[-0.876]~[0.000]\n", "\n", "Prismatic joint information. n_pri_joint:[0]\n", "\n", "Control information. n_ctrl:[7]\n", " [0/7] [actuator_joint1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [1/7] [actuator_joint2] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [2/7] [actuator_joint3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [3/7] [actuator_joint4] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [4/7] [actuator_joint5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [5/7] [actuator_joint6] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [6/7] [actuator_finger_joint] range:[0.000]~[0.725] gear:[1.00] type:[JOINT]\n", "\n", "Camera information. n_cam:[4]\n", " [0/4] [agentview] fov:[60.0]\n", " [1/4] [topview] fov:[90.0]\n", " [2/4] [sideview] fov:[90.0]\n", " [3/4] [egocentric] fov:[90.0]\n", "\n", "n_sensor:[0]\n", "sensor_names:[]\n", "n_site:[6]\n", "site_names:['bottom_site_mug_5', 'top_site_mug_5', 'horizontal_radius_site_mug_5', 'bottom_site_plate_11', 'top_site_plate_11', 'horizontal_radius_site_plate_11']\n", "-----------------------------------------------------------------------------\n", "env:[Tabletop] reset\n", "env:[Tabletop] reset\n", "env:[Tabletop] initalize viewer\n", "DONE INITIALIZATION\n" ] } ], "source": [ "from mujoco_env.y_env import SimpleEnv\n", "# xml_path = './asset/example_scene_y.xml'\n", "xml_path = './asset/example_scene_nova5.xml'\n", "# 这里选择具体场景即可\n", "PnPEnv = SimpleEnv(xml_path, action_type='joint_angle')" ] }, { "cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "DONE INITIALIZATION\n", "Replay finished. steps=642, episode_len=642\n" ] } ], "source": [ "# Replay mode options:\n", "# - OVERLAY_SOURCE = \"dataset\": top-right/bottom-right show recorded images.\n", "# - OVERLAY_SOURCE = \"sim\": top-right/bottom-right show current simulation camera images.\n", "OVERLAY_SOURCE = \"sim\" # \"sim\" or \"dataset\"\n", "RENDER_HZ = 60\n", "ACTION_HZ = 20\n", "\n", "step = 0\n", "iter_dataloader = iter(dataloader)\n", "PnPEnv.reset()\n", "\n", "action = np.zeros(7, dtype=np.float32)\n", "done = False\n", "\n", "while PnPEnv.env.is_viewer_alive() and not done:\n", " # Update action from dataset at ACTION_HZ\n", " update_action = PnPEnv.env.loop_every(HZ=ACTION_HZ)\n", " if update_action:\n", " try:\n", " data = next(iter_dataloader)\n", " except StopIteration:\n", " break\n", "\n", " if step == 0:\n", " # Reset object pose from the first frame in this episode\n", " PnPEnv.set_obj_pose(data['obj_init'][0, :3], data['obj_init'][0, 3:])\n", "\n", " # LeRobot action shape in this notebook: [B=1, 7]\n", " action = data['action'][0].numpy().astype(np.float32)\n", " _ = PnPEnv.step(action)\n", " step += 1\n", "\n", " if OVERLAY_SOURCE == \"dataset\":\n", " rgb_agent = (data['observation.image'][0].numpy() * 255).astype(np.uint8)\n", " rgb_ego = (data['observation.wrist_image'][0].numpy() * 255).astype(np.uint8)\n", " # CHW -> HWC\n", " PnPEnv.rgb_agent = np.transpose(rgb_agent, (1, 2, 0))\n", " PnPEnv.rgb_ego = np.transpose(rgb_ego, (1, 2, 0))\n", " PnPEnv.rgb_side = np.zeros((480, 640, 3), dtype=np.uint8)\n", "\n", " if step >= len(episode_sampler):\n", " done = True\n", "\n", " # Advance physics every tick (same pattern as collect loop)\n", " PnPEnv.step_env()\n", "\n", " # Keep viewer responsive. If using sim overlay, refresh camera frames each render tick.\n", " if PnPEnv.env.loop_every(HZ=RENDER_HZ):\n", " if OVERLAY_SOURCE == \"sim\":\n", " PnPEnv.grab_image()\n", " PnPEnv.render(teleop=True)\n", "\n", "print(f\"Replay finished. steps={step}, episode_len={len(episode_sampler)}\")\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 29, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1 642\n" ] } ], "source": [ "print(dataset.num_episodes, dataset.num_frames)\n" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [], "source": [ "# 为了减少卡顿，建议确保系统内存占用充裕，可关闭不必要的后台程序，或考虑使用远程 Linux 环境以减轻 Windows 负担。" ] }, { "cell_type": "code", "execution_count": 31, "metadata": {}, "outputs": [], "source": [ "PnPEnv.env.close_viewer()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 【可选】为其他版本保存 Stats.json\n" ] }, { "cell_type": "code", "execution_count": 32, "metadata": {}, "outputs": [], "source": [ "stats = dataset.meta.stats\n", "PATH = dataset.root / 'meta' / 'stats.json'\n", "stats = serialize_dict(stats)\n", "\n", "write_json(stats, PATH)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] } ], "metadata": { "kernelspec": { "display_name": "py310", "language": "python", "name": "py310" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.19" } }, "nbformat": 4, "nbformat_minor": 4 }