Add visualizations to demos.

Author: kylc

demo/01_intro.py

@@ -3,6 +3,7 @@
 from math import sqrt
 
 import numpy as np
+import rerun as rr
 import planning_kit as pkit
 from planning_kit import StateSpace, Constraint
 
@@ -13,21 +14,38 @@
 # We define lower and upper boundaries of the space in R^3.
 euclidean = StateSpace.euclidean(-np.ones(3), np.ones(3))
 
+R = 0.95
 
 # A sample is valid if it lies outside a sphere of radius 0.95.
 def is_valid(s):
-    return sqrt(s[0] ** 2 + s[1] ** 2 + s[2] ** 3) > 0.95
+    return sqrt(s[0] ** 2 + s[1] ** 2 + s[2] ** 3) > R
 
 
 # Now we can use a variety of the built-in sampling-based motion planning
 # algorithms.
+start = -np.ones(3)
+goal = np.ones(3)
 tree = pkit.rrt_connect(
     euclidean,  # the search space
     is_valid,  # is a particular sample valid?
-    start=-np.ones(3),  # the start state
-    goal=np.ones(3),  # the goal state
+    start=start,  # the start state
+    goal=goal,  # the goal state
     discretization=0.01,  # motion validation resolution
     # planner-specific parameters
     steering_dist=0.1,
     n=10000,
 )
+
+# Log everything to Rerun!
+rr.init("pk.01_intro", spawn=True)
+rr.log_view_coordinates("space", xyz="FLU", timeless=True)
+rr.log_points("space/rrt_tree", positions=[p for p in tree.points()])
+rr.log_line_strip(
+    "space/shortest_path",
+    positions=tree.shortest_path(euclidean, start, goal).points(),
+)
+
+# Sample the surface of the sphere so we can visualize the keep-out region.
+xs = np.random.uniform(low=-1.0, high=1.0, size=(1000, 3))
+xs = R * xs / np.linalg.norm(xs, axis=1)[:, None]
+rr.log_points("space/sphere", positions=xs)

demo/02_custom_space.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import numpy as np
+import rerun as rr
 from planning_kit import StateSpace
 
 
@@ -24,3 +25,8 @@ def interpolate(self, a, b, alpha):
 # Once we wrap our custom implementation in a StateSpace object, we can use it
 # as any other.
 space = StateSpace.custom(MyEuclideanSpace(dim=3))
+
+# Log everything to Rerun!
+rr.init("pk.02_custom_space", spawn=True)
+rr.log_view_coordinates("space", xyz="FLU", timeless=True)
+rr.log_points("space", positions=[space.sample_uniform() for _ in range(0, 1000)])

demo/03_constraints.py

@@ -1,14 +1,33 @@
 #!/usr/bin/env python3
 
 import numpy as np
+import rerun as rr
 from planning_kit import Constraint, StateSpace
 
 
-def sphere(s):
-    return [s[0] ** 2 + s[1] ** 2 + s[2] ** 2 - 1.0]
+def sinusoid_xy_plane(amp=1.0, freq=1.0):
+    def inner(p):
+        s = np.sin(2.0 * np.pi * freq * p[0])
+        c = np.cos(2.0 * np.pi * freq * p[1])
+        z = amp * (s + c)
+        return [p[2] - z]
 
+    return inner
 
-sphere_constraint = Constraint(3, 1, sphere)
+
+constraint = Constraint(3, 1, sinusoid_xy_plane(amp=0.1, freq=0.707))
 
 ambient = StateSpace.euclidean(-np.ones(3), np.ones(3))
-projected = StateSpace.projected(ambient, sphere_constraint)
+projected = StateSpace.projected(ambient, constraint)
+
+# Log everything to Rerun!
+rr.init("pk.03_constraints", spawn=True)
+rr.log_view_coordinates("space", xyz="FLU", timeless=True)
+rr.log_points(
+    "space/ambient",
+    positions=[ambient.sample_uniform() for _ in range(0, 500)],
+    colors=[0.3, 0.3, 0.3],
+)
+rr.log_points(
+    "space/constrained", positions=[projected.sample_uniform() for _ in range(0, 2000)]
+)