Function main

examples/09_mujoco_streaming.py:13–86

(args)

def main(args):
    import mujoco

    # Load the scene
    xml_path = str(ASSETS_DIR / "scenes" / "franka_emika_panda" / "scene_blockpush.xml")
    model = mujoco.MjModel.from_xml_path(xml_path)
    data = mujoco.MjData(model)

    from avp_stream import VisionProStreamer
    streamer = VisionProStreamer(ip=args.ip, record=False)

    # attach_to format: [x, y, z, yaw_degrees]
    attach_to = [0.2, 1.0, 0.7, -90]

    streamer.configure_mujoco(
        xml_path=xml_path,
        model=model,
        data=data,
        relative_to=attach_to,
        grpc_port=args.port,
    )

    streamer.start_webrtc()
    
    # Wait for connection and USDZ transfer
    print("[CROSS-NETWORK] Waiting for WebRTC connection and scene transfer...")
    if not streamer.wait_for_sim_channel(timeout=60.0):
        print("Timeout or error waiting for connection!")
        return


    logs_dir = ASSETS_DIR / "logs"
    episodes = sorted(logs_dir.glob("ep*.npz"))

    try:
        for ep_idx in range(1, min(5, len(episodes) + 1)):
            traj_path = logs_dir / f"ep{ep_idx}.npz"
            if not traj_path.exists():
                continue

            print(f"📼 Playing episode {ep_idx}...")
            traj = np.load(traj_path)

            qpos_log = traj["qpos"]
            qvel_log = traj["qvel"]
            ctrl_log = traj["ctrl"]
            mocap_log = traj["mocap"]

            T = qpos_log.shape[0]

            # Initialize state
            data.qpos[:] = qpos_log[0]
            data.qvel[:] = qvel_log[0]

            if mocap_log.shape[0] > 0:
                data.mocap_pos[1] = mocap_log[0, :3]
                data.mocap_quat[1] = mocap_log[0, 3:]