Online PlanningΒΆ

Run online planning in collision aware environments.

All examples can be run by first cloning the PyRoki repository, which includes the pyroki_snippets implementation details.

  1import time
  2import viser
  3import numpy as np
  4from robot_descriptions.loaders.yourdfpy import load_robot_description
  5
  6import pyroki as pk
  7from pyroki.collision import HalfSpace, RobotCollision, Sphere
  8from viser.extras import ViserUrdf
  9import pyroki_snippets as pks
 10
 11
 12def main():
 13    """Main function for online planning with collision."""
 14    urdf = load_robot_description("panda_description")
 15    target_link_name = "panda_hand"
 16    robot = pk.Robot.from_urdf(urdf)
 17
 18    robot_coll = RobotCollision.from_urdf(urdf)
 19    plane_coll = HalfSpace.from_point_and_normal(
 20        np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])
 21    )
 22    sphere_coll = Sphere.from_center_and_radius(
 23        np.array([0.0, 0.0, 0.0]), np.array([0.05])
 24    )
 25
 26    # Define the online planning parameters.
 27    len_traj, dt = 5, 0.1
 28
 29    # Set up visualizer.
 30    server = viser.ViserServer()
 31    server.scene.add_grid("/ground", width=2, height=2, cell_size=0.1)
 32    urdf_vis = ViserUrdf(server, urdf, root_node_name="/robot")
 33
 34    # Create interactive controller for IK target.
 35    ik_target_handle = server.scene.add_transform_controls(
 36        "/ik_target", scale=0.2, position=(0.3, 0.0, 0.5), wxyz=(0, 0, 1, 0)
 37    )
 38
 39    # Create interactive controller and mesh for the sphere obstacle.
 40    sphere_handle = server.scene.add_transform_controls(
 41        "/obstacle", scale=0.2, position=(0.4, 0.3, 0.4)
 42    )
 43    server.scene.add_mesh_trimesh("/obstacle/mesh", mesh=sphere_coll.to_trimesh())
 44    target_frame_handle = server.scene.add_batched_axes(
 45        "target_frame",
 46        axes_length=0.05,
 47        axes_radius=0.005,
 48        batched_positions=np.zeros((25, 3)),
 49        batched_wxyzs=np.array([[1.0, 0.0, 0.0, 0.0]] * 25),
 50    )
 51
 52    timing_handle = server.gui.add_number("Elapsed (ms)", 0.001, disabled=True)
 53
 54    sol_pos, sol_wxyz = None, None
 55    sol_traj = np.array(
 56        robot.joint_var_cls.default_factory()[None].repeat(len_traj, axis=0)
 57    )
 58    while True:
 59        start_time = time.time()
 60
 61        sphere_coll_world_current = sphere_coll.transform_from_pos_wxyz(
 62            position=np.array(sphere_handle.position),
 63            wxyz=np.array(sphere_handle.wxyz),
 64        )
 65
 66        world_coll_list = [plane_coll, sphere_coll_world_current]
 67        sol_traj, sol_pos, sol_wxyz = pks.solve_online_planning(
 68            robot=robot,
 69            robot_coll=robot_coll,
 70            world_coll=world_coll_list,
 71            target_link_name=target_link_name,
 72            target_position=np.array(ik_target_handle.position),
 73            target_wxyz=np.array(ik_target_handle.wxyz),
 74            timesteps=len_traj,
 75            dt=dt,
 76            start_cfg=sol_traj[0],
 77            prev_sols=sol_traj,
 78        )
 79
 80        # Update timing handle.
 81        timing_handle.value = (
 82            0.99 * timing_handle.value + 0.01 * (time.time() - start_time) * 1000
 83        )
 84
 85        # Update visualizer.
 86        urdf_vis.update_cfg(
 87            sol_traj[0]
 88        )  # The first step of the online trajectory solution.
 89
 90        # Update the planned trajectory visualization.
 91        if hasattr(target_frame_handle, "batched_positions"):
 92            target_frame_handle.batched_positions = np.array(sol_pos)  # type: ignore[attr-defined]
 93            target_frame_handle.batched_wxyzs = np.array(sol_wxyz)  # type: ignore[attr-defined]
 94        else:
 95            # This is an older version of Viser.
 96            target_frame_handle.positions_batched = np.array(sol_pos)  # type: ignore[attr-defined]
 97            target_frame_handle.wxyzs_batched = np.array(sol_wxyz)  # type: ignore[attr-defined]
 98
 99
100if __name__ == "__main__":
101    main()