{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 可视化你的数据\n", "\n", "\n", "\n", "基于重建后的仿真场景可视化你的动作。\n", "\n", "主窗口会回放动作。\n", "\n", "右上和右下叠加图像来自数据集。\n" ] }, { "cell_type": "code", "execution_count": 1, "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": "9750632bb0f94d5e82515c8263a6ba12", "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", "REPO_NAME = 'datawhale_eai_pnp_xarm7' #pick-and-place\n", "NUM_DEMO = 1 # Number of demonstrations to collect\n", "ROOT = \"./demo_data_xarm7\" # The root directory to save the demonstrations\n", "dataset = LeRobotDataset('datawhale_eai_pnp_xarm7', root=\"./demo_data_xarm7\") # 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": 2, "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": 3, "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": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Env Info] 运行环境已就绪，当前工作目录: C:\\Users\\kewei\\lerobot-mujoco-tutorial\n" ] } ], "source": [ "import env_config" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "-----------------------------------------------------------------------------\n", "name:[Tabletop] dt:[0.002] HZ:[500]\n", " n_qpos:[27] n_qvel:[25] n_qacc:[25] n_ctrl:[8]\n", " integrator:[IMPLICITFAST]\n", "\n", "n_body:[25]\n", " [0/25] [world] mass:[0.00]kg\n", " [1/25] [front_object_table] mass:[1.00]kg\n", " [2/25] [camera] mass:[0.00]kg\n", " [3/25] [camera2] mass:[0.00]kg\n", " [4/25] [camera3] mass:[0.00]kg\n", " [5/25] [link_base] mass:[0.89]kg\n", " [6/25] [link1] mass:[2.38]kg\n", " [7/25] [link2] mass:[1.87]kg\n", " [8/25] [link3] mass:[1.64]kg\n", " [9/25] [link4] mass:[1.73]kg\n", " [10/25] [link5] mass:[1.32]kg\n", " [11/25] [link6] mass:[1.32]kg\n", " [12/25] [link7] mass:[0.17]kg\n", " [13/25] [xarm_gripper_base_link] mass:[0.54]kg\n", " [14/25] [left_outer_knuckle] mass:[0.03]kg\n", " [15/25] [left_finger] mass:[0.05]kg\n", " [16/25] [left_inner_knuckle] mass:[0.02]kg\n", " [17/25] [right_outer_knuckle] mass:[0.03]kg\n", " [18/25] [right_finger] mass:[0.05]kg\n", " [19/25] [right_inner_knuckle] mass:[0.02]kg\n", " [20/25] [camera_wrist] mass:[0.00]kg\n", " [21/25] [body_obj_mug_5] mass:[0.00]kg\n", " [22/25] [object_mug_5] mass:[0.08]kg\n", " [23/25] [body_obj_plate_11] mass:[0.00]kg\n", " [24/25] [object_plate_11] mass:[0.10]kg\n", "body_total_mass:[13.25]kg\n", "\n", "n_geom:[87]\n", "geom_names:['floor', 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, 'left_finger_pad_1', 'left_finger_pad_2', None, None, None, 'right_finger_pad_1', 'right_finger_pad_2', 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:[82]\n", "mesh_names:['link_base', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'link7', 'end_tool', 'base_link', 'left_outer_knuckle', 'left_finger', 'left_inner_knuckle', 'right_outer_knuckle', 'right_finger', 'right_inner_knuckle', '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:[15]\n", " [0/15] [joint1] axis:[0. 0. 1.]\n", " [1/15] [joint2] axis:[0. 0. 1.]\n", " [2/15] [joint3] axis:[0. 0. 1.]\n", " [3/15] [joint4] axis:[0. 0. 1.]\n", " [4/15] [joint5] axis:[0. 0. 1.]\n", " [5/15] [joint6] axis:[0. 0. 1.]\n", " [6/15] [joint7] axis:[0. 0. 1.]\n", " [7/15] [left_driver_joint] axis:[1. 0. 0.]\n", " [8/15] [left_finger_joint] axis:[-1. 0. 0.]\n", " [9/15] [left_inner_knuckle_joint] axis:[1. 0. 0.]\n", " [10/15] [right_driver_joint] axis:[-1. 0. 0.]\n", " [11/15] [right_finger_joint] axis:[1. 0. 0.]\n", " [12/15] [right_inner_knuckle_joint] axis:[-1. 0. 0.]\n", " [13/15] [None] axis:[0. 0. 1.]\n", " [14/15] [None] axis:[0. 0. 1.]\n", "\n", "n_dof:[25] (=number of rows of Jacobian)\n", " [0/25] [None] attached joint:[joint1] body:[link1]\n", " [1/25] [None] attached joint:[joint2] body:[link2]\n", " [2/25] [None] attached joint:[joint3] body:[link3]\n", " [3/25] [None] attached joint:[joint4] body:[link4]\n", " [4/25] [None] attached joint:[joint5] body:[link5]\n", " [5/25] [None] attached joint:[joint6] body:[link6]\n", " [6/25] [None] attached joint:[joint7] body:[link7]\n", " [7/25] [None] attached joint:[left_driver_joint] body:[left_outer_knuckle]\n", " [8/25] [None] attached joint:[left_finger_joint] body:[left_finger]\n", " [9/25] [None] attached joint:[left_inner_knuckle_joint] body:[left_inner_knuckle]\n", " [10/25] [None] attached joint:[right_driver_joint] body:[right_outer_knuckle]\n", " [11/25] [None] attached joint:[right_finger_joint] body:[right_finger]\n", " [12/25] [None] attached joint:[right_inner_knuckle_joint] body:[right_inner_knuckle]\n", " [13/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [14/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [15/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [16/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [17/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [18/25] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [19/25] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [20/25] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [21/25] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [22/25] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [23/25] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [24/25] [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:[13]\n", " [0/13] [joint1] range:[-6.283]~[6.283]\n", " [1/13] [joint2] range:[-2.059]~[2.094]\n", " [2/13] [joint3] range:[-6.283]~[6.283]\n", " [3/13] [joint4] range:[-0.192]~[3.927]\n", " [4/13] [joint5] range:[-6.283]~[6.283]\n", " [5/13] [joint6] range:[-1.693]~[3.142]\n", " [6/13] [joint7] range:[-6.283]~[6.283]\n", " [7/13] [left_driver_joint] range:[0.000]~[0.850]\n", " [8/13] [left_finger_joint] range:[0.000]~[0.850]\n", " [9/13] [left_inner_knuckle_joint] range:[0.000]~[0.850]\n", " [10/13] [right_driver_joint] range:[0.000]~[0.850]\n", " [11/13] [right_finger_joint] range:[0.000]~[0.850]\n", " [12/13] [right_inner_knuckle_joint] range:[0.000]~[0.850]\n", "\n", "Prismatic joint information. n_pri_joint:[0]\n", "\n", "Control information. n_ctrl:[8]\n", " [0/8] [act1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [1/8] [act2] range:[-2.059]~[2.094] gear:[1.00] type:[JOINT]\n", " [2/8] [act3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [3/8] [act4] range:[-0.192]~[3.927] gear:[1.00] type:[JOINT]\n", " [4/8] [act5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [5/8] [act6] range:[-1.693]~[3.142] gear:[1.00] type:[JOINT]\n", " [6/8] [act7] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [7/8] [gripper] range:[0.000]~[255.000] gear:[1.00] type:[TENDON]\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:[7]\n", "site_names:['link_tcp', '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", "[IK] Using pybullet IK backend. ee_link=xarm_gripper_base_link idx=9\n" ] } ], "source": [ "from mujoco_env.y_env import SimpleEnv\n", "# xml_path = './asset/example_scene_y.xml'\n", "xml_path = './asset/example_scene_xarm7.xml'\n", "# 这里选择具体场景即可\n", "PnPEnv = SimpleEnv(xml_path, action_type='joint_angle')" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "DONE INITIALIZATION\n", "Replay finished. steps=827, episode_len=827\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": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1 827\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 }