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