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()