openscvx 0.4.1.dev170__tar.gz → 0.4.1.dev172__tar.gz

{openscvx-0.4.1.dev170 → openscvx-0.4.1.dev172}/.gitignore

@@ -27,6 +27,9 @@ solver/*
 # Images
 *.eps
 
+# Animation renders
+examples/animations/mp4/
+
 # Numpy files
 *.npy
 

{openscvx-0.4.1.dev170/openscvx.egg-info → openscvx-0.4.1.dev172}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.4.1.dev170
+Version: 0.4.1.dev172
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-0.4.1.dev170 → openscvx-0.4.1.dev172}/docs/UsersGuide/07_lie.md

@@ -364,6 +364,98 @@ results = problem.post_process()
 
 After solving, `results.trajectory["T_ee"]` gives you the dense end-effector transform along the propagated trajectory, ready to feed straight into a viser animation or a downstream controller.
 
+## Extension: Self-Collision and Obstacle Avoidance
+
+The same intermediate `joint_transforms` we snapshotted above give us everything we need to constrain the *arm itself*, not just its end-effector. The [7-DOF arm with collisions](../Examples/arm/7_dof_arm_collision.md) example extends the pick-and-place setup with two extra constraints:
+
+- **Self-collision** between every pair of non-adjacent links.
+- **Obstacle avoidance** between every link and a spherical obstacle placed on the table.
+
+The idea is to approximate each link as a **capsule** (a sphere-swept line segment) between two consecutive keypoints, sample each capsule at a small grid of parameters $t \in [0, 1]$, and enforce a pairwise distance bound at every sample. `ox.Vmap` turns the inner loop into a single batched evaluation.
+
+First, lift each keypoint from its home position into the world frame using the transforms we already built. We pad with a trailing `1` so the $4 \times 4$ homogeneous transform acts on it correctly:
+
+```python
+keypoint_home_pos = {
+    "base":     np.array([0.0,          0.0, 0.0, 1.0]),
+    "shoulder": np.array([0.0,          0.0, d1,  1.0]),
+    "elbow":    np.array([a2,           0.0, d1,  1.0]),
+    "wrist":    np.array([a2 + a3,      0.0, d1,  1.0]),
+    "ee":       np.array([a2 + a3 + a4, 0.0, d1,  1.0]),
+}
+
+def _keypoint_world(name):
+    p_hom = joint_transforms[name] @ ox.Constant(keypoint_home_pos[name])
+    return ox.Concat(p_hom[0], p_hom[1], p_hom[2])
+
+kp = {name: _keypoint_world(name) for name in joint_names}
+```
+
+Then describe each link as `(start, end, radius)`:
+
+```python
+link_specs = [
+    ("base",     "shoulder", 0.06),
+    ("shoulder", "elbow",    0.05),
+    ("elbow",    "wrist",    0.04),
+    ("wrist",    "ee",       0.035),
+]
+```
+
+### Self-collision
+
+For every pair of non-adjacent capsules, sample both segments on a $t \times t$ grid and enforce that the sum of radii stays under the sampled distance:
+
+```python
+n_samples = 4
+ts = np.linspace(0.0, 1.0, n_samples)
+tt = np.stack(np.meshgrid(ts, ts, indexing="ij"), axis=-1).reshape(-1, 2)
+
+for i in range(len(link_specs)):
+    for j in range(i + 2, len(link_specs)):  # skip adjacent (shared joint)
+        a_start, a_end, r_a = link_specs[i]
+        b_start, b_end, r_b = link_specs[j]
+        pa0, pa1 = kp[a_start], kp[a_end]
+        pb0, pb1 = kp[b_start], kp[b_end]
+        r_sum = r_a + r_b
+
+        dists = ox.Vmap(
+            lambda t, pa0=pa0, pa1=pa1, pb0=pb0, pb1=pb1: ox.linalg.Norm(
+                ((1 - t[0]) * pa0 + t[0] * pa1)
+                - ((1 - t[1]) * pb0 + t[1] * pb1)
+            ),
+            batch=tt,
+        )
+        constraints.append(ox.ctcs(r_sum <= dists))
+```
+
+Skipping `j = i + 1` matters: adjacent capsules share a joint, so their minimum distance is identically zero and enforcing a strict capsule-capsule bound between them would be infeasible by construction.
+
+### Obstacle avoidance
+
+The same pattern handles a spherical obstacle — sample every link and keep it outside the inflated sphere:
+
+```python
+obstacle_center = np.array([0.35, 0.0, 0.25])
+obstacle_radius = 0.06
+
+for a_start, a_end, r_link in link_specs:
+    pa0, pa1 = kp[a_start], kp[a_end]
+    r_clear = obstacle_radius + r_link
+    obs_dists = ox.Vmap(
+        lambda t, pa0=pa0, pa1=pa1: ox.linalg.Norm(
+            ((1 - t) * pa0 + t * pa1) - ox.Constant(obstacle_center)
+        ),
+        batch=ts,
+    )
+    constraints.append(ox.ctcs(r_clear <= obs_dists))
+```
+
+Because each bound is wrapped in `ox.ctcs`, distances are enforced **continuously** along the trajectory — capsules cannot tunnel between nodes.
+
+!!! warning "Numerical performance"
+    Solve times on this example are currently noticeably slower than we would like when self-collision checking is included; expect a meaningful jump compared to the raw pick-and-place. We are actively working on it.
+
 ## Key Takeaways
 
 1. **Start simple with the minimal representation.** Joint angles + velocities give a pure ODE that drops cleanly into OpenSCvx's modeling stack — the easiest place to begin. Maximal coordinates are also a strong (and in some respects superior) choice for trajopt and will get their own tutorial in the future. The consequence of working in joint space is that task-space quantities are no longer states — they are symbolic functions of the state.
@@ -377,4 +469,5 @@ After solving, `results.trajectory["T_ee"]` gives you the dense end-effector tra
 - [API Reference: Lie algebra](../Reference/symbolic/expr/lie/index.md)
 - [Three-link arm Example](../Examples/arm/3_dof_arm.md)
 - [Seven-link arm Example](../Examples/arm/7_dof_arm.md)
+- [Seven-link arm with self-collision Example](../Examples/arm/7_dof_arm_collision.md)
 - [Seven-link arm with viewcone Example](../Examples/arm/7_dof_arm_vp.md)

openscvx-0.4.1.dev172/examples/animations/7_dof_arm.py

@@ -0,0 +1,225 @@
+"""Cinematic offline render of the 7-DOF arm pick-and-place example.
+
+The trajectory optimization problem itself lives in
+``examples/arm/7_dof_arm.py``; this file imports that ``problem`` and the
+robot parameters, solves it, builds a viser scene with arm links / joint
+frames / EE trail, and drives it frame-by-frame while piping raw RGB into
+ffmpeg to produce an mp4.
+
+Run it with::
+
+    python examples/animations/7_dof_arm.py
+
+The script prints a viser URL and waits. Open the URL in a browser — as soon
+as the client connects, the render begins. Requires ``ffmpeg`` on ``PATH``;
+``openscvx`` does not depend on it.
+
+Only one camera mode is provided — ``"overview"`` — a static elevated camera
+that frames the full pick-and-place workspace.
+
+Tweak ``OUTPUT_PATH`` / ``WIDTH`` / ``HEIGHT`` / ``FPS`` below for different
+output variants.
+"""
+
+import importlib
+import os
+import sys
+
+import numpy as np
+
+# Add the project root so `examples.*` imports resolve.
+current_dir = os.path.dirname(os.path.abspath(__file__))
+grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
+sys.path.append(grandparent_dir)
+
+from examples.animations._camera import overview_pose
+from examples.animations._render import render_animation_to_video
+from examples.plotting_viser import AnimatedServerHandle
+from openscvx.plotting.viser import (
+    add_animated_trail,
+    add_ellipsoid_obstacles,
+    add_ghost_trajectory,
+    add_position_marker,
+    add_target_markers,
+    compute_velocity_colors,
+    create_server,
+)
+
+# Import the arm example by path (filename starts with a digit).
+_arm_spec = importlib.util.spec_from_file_location(
+    "arm_7dof", os.path.join(grandparent_dir, "examples", "arm", "7_dof_arm_collision.py")
+)
+_arm_mod = importlib.util.module_from_spec(_arm_spec)
+_arm_spec.loader.exec_module(_arm_mod)
+
+problem = _arm_mod.problem
+d1 = _arm_mod.d1
+a2 = _arm_mod.a2
+a3 = _arm_mod.a3
+a4 = _arm_mod.a4
+joint_names = _arm_mod.joint_names
+waypoint_names = _arm_mod.waypoint_names
+waypoint_positions = _arm_mod.waypoint_positions
+obstacle_center = _arm_mod.obstacle_center
+obstacle_radius = _arm_mod.obstacle_radius
+
+# Keypoint home positions (derived from link lengths — defined in __main__ of
+# the arm example, so we replicate them here).
+keypoint_home_positions = {
+    "base": np.array([0.0, 0.0, 0.0]),
+    "shoulder": np.array([0.0, 0.0, d1]),
+    "elbow": np.array([a2, 0.0, d1]),
+    "wrist": np.array([a2 + a3, 0.0, d1]),
+    "ee": np.array([a2 + a3 + a4, 0.0, d1]),
+}
+
+# --- Render settings ---------------------------------------------------------
+OUTPUT_PATH = os.path.join(current_dir, "mp4", "7_dof_arm_overview.mp4")
+WIDTH = 1080
+HEIGHT = 1080
+FPS = 60
+CRF = 16
+
+# Oversampling for smooth trails.
+STRIDE = 4
+PROPAGATION_HZ = FPS * STRIDE
+
+# --- Camera settings ---------------------------------------------------------
+OVERVIEW_AZIMUTH = np.radians(60.0)
+OVERVIEW_ELEVATION = np.radians(15.0)
+OVERVIEW_RADIUS_MARGIN = 1.0
+OVERVIEW_FOV_DEG = 60.0
+
+
+if __name__ == "__main__":
+    import jaxlie
+
+    problem.settings.prp.dt = 1.0 / PROPAGATION_HZ
+    problem.initialize()
+    problem.solve()
+    results = problem.post_process()
+
+    ee_pos = np.asarray(results.trajectory["ee_position"], dtype=np.float64)
+    n_frames = len(ee_pos)
+
+    # -- Extract joint keypoint positions from propagated transforms -----------
+    keypoints = np.zeros((n_frames, 5, 3))
+    joint_quats = np.zeros((n_frames, 5, 4))
+    for k, name in enumerate(joint_names):
+        T_k = np.asarray(results.trajectory[f"T_{name}"])  # (T, 4, 4)
+        p_home = np.append(keypoint_home_positions[name], 1.0)
+        keypoints[:, k] = (T_k @ p_home)[:, :3]
+        joint_quats[:, k] = np.array(
+            [jaxlie.SO3.from_matrix(T_k[t, :3, :3]).wxyz for t in range(n_frames)]
+        )
+
+    # -- Build viser scene (mirrors examples/arm/7_dof_arm.py) -----------------
+    server = create_server(ee_pos, show_grid=False)
+    server.scene.add_grid("/grid", width=1.5, height=1.5, cell_size=0.25)
+    server.scene.add_frame("/origin", axes_length=0.1, axes_radius=0.003)
+
+    # Waypoint markers
+    marker_colors = {
+        "home": (100, 150, 255),
+        "pre_grasp": (255, 180, 50),
+        "grasp": (255, 50, 50),
+        "pre_place": (255, 180, 50),
+        "place": (255, 50, 50),
+    }
+    add_target_markers(
+        server,
+        waypoint_positions,
+        radius=0.015,
+        colors=[marker_colors[name] for name in waypoint_names],
+    )
+
+    add_ellipsoid_obstacles(
+        server,
+        centers=[obstacle_center],
+        radii=[np.full(3, 1.0 / obstacle_radius)],
+    )
+
+    # Ghost EE trajectory + animated trail
+    ee_colors = compute_velocity_colors(np.asarray(results.trajectory.get("velocity")))
+    add_ghost_trajectory(server, ee_pos, ee_colors, point_size=0.005)
+    _, update_trail = add_animated_trail(server, ee_pos, ee_colors, point_size=0.008)
+
+    # Animated EE position marker
+    _, update_marker = add_position_marker(server, ee_pos, radius=0.015)
+
+    # Animated arm links (line segments between consecutive keypoints)
+    link_rgb = np.array(
+        [
+            [180, 180, 180],  # base -> shoulder
+            [100, 180, 255],  # shoulder -> elbow
+            [100, 255, 150],  # elbow -> wrist
+            [255, 200, 100],  # wrist -> ee
+        ],
+        dtype=np.uint8,
+    )
+    link_colors = np.stack([link_rgb, link_rgb], axis=1)  # (4, 2, 3)
+
+    init_points = np.stack(
+        [np.stack([keypoints[0, k], keypoints[0, k + 1]]) for k in range(4)]
+    ).astype(np.float32)
+
+    arm_handle = server.scene.add_line_segments(
+        "/arm_links",
+        points=init_points,
+        colors=link_colors,
+        line_width=5.0,
+    )
+
+    # Joint coordinate frames
+    joint_frame_handles = []
+    for k in range(5):
+        h = server.scene.add_frame(
+            f"/joint_{joint_names[k]}",
+            wxyz=joint_quats[0, k].astype(np.float32),
+            position=keypoints[0, k].astype(np.float32),
+            axes_length=0.06,
+            axes_radius=0.002,
+        )
+        joint_frame_handles.append(h)
+
+    def update_arm(frame_idx: int) -> None:
+        pts = np.stack(
+            [np.stack([keypoints[frame_idx, k], keypoints[frame_idx, k + 1]]) for k in range(4)]
+        ).astype(np.float32)
+        arm_handle.points = pts
+        for k, h in enumerate(joint_frame_handles):
+            h.position = keypoints[frame_idx, k].astype(np.float32)
+            h.wxyz = joint_quats[frame_idx, k].astype(np.float32)
+
+    # -- Manual-step handle for the renderer -----------------------------------
+    traj_time = np.asarray(results.trajectory["time"], dtype=np.float64).flatten()
+    handle = AnimatedServerHandle(
+        server=server,
+        traj_time=traj_time,
+        update_callbacks=[update_trail, update_marker, update_arm],
+    )
+
+    # -- Camera: overview framing the full workspace ---------------------------
+    # Frame on waypoint positions + EE path so the full motion is visible.
+    all_points = np.vstack([ee_pos, np.array(waypoint_positions)])
+    static_pose = overview_pose(
+        all_points,
+        azimuth=OVERVIEW_AZIMUTH,
+        elevation=OVERVIEW_ELEVATION,
+        radius_margin=OVERVIEW_RADIUS_MARGIN,
+        fov_deg=OVERVIEW_FOV_DEG,
+    )
+
+    def camera_pose_fn(frame_idx: int):
+        return static_pose
+
+    render_animation_to_video(
+        handle,
+        OUTPUT_PATH,
+        camera_pose_fn,
+        width=WIDTH,
+        height=HEIGHT,
+        fps=FPS,
+        crf=CRF,
+        stride=STRIDE,
+    )

openscvx-0.4.1.dev172/examples/animations/_camera.py

@@ -0,0 +1,142 @@
+"""Reusable camera pose helpers for viser animation rendering.
+
+All functions return ``(cam_position, cam_wxyz, look_at)`` tuples suitable for
+passing to ``render_animation_to_video``'s ``camera_pose_fn`` argument. They
+depend only on numpy and ``viser.transforms`` — no openscvx imports.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import viser.transforms as vtf
+
+
+def look_at_wxyz(pos: np.ndarray, target: np.ndarray, up: np.ndarray) -> np.ndarray:
+    """Quaternion (w,x,y,z) for a camera at ``pos`` looking at ``target``.
+
+    Uses the OpenCV camera convention that viser expects: +X right, +Y down,
+    +Z forward. See ``examples/animations/camera_control_notes.md``.
+    """
+    forward = target - pos
+    forward /= np.linalg.norm(forward)
+    right = np.cross(forward, up)
+    right_norm = np.linalg.norm(right)
+    if right_norm < 1e-6:
+        # Gimbal lock: forward is (nearly) parallel to up. Pick an arbitrary
+        # world axis that isn't, so the camera stays defined. The chosen axis
+        # determines the "roll" of the camera at the singularity — not great
+        # cinematically, but prevents a NaN crash.
+        fallback = np.array([1.0, 0.0, 0.0]) if abs(forward[0]) < 0.9 else np.array([0.0, 1.0, 0.0])
+        right = np.cross(forward, fallback)
+        right_norm = np.linalg.norm(right)
+    right /= right_norm
+    cam_down = np.cross(forward, right)  # = -world_up projected perp to forward
+    R_world_cam = np.stack([right, cam_down, forward], axis=1)
+    return vtf.SO3.from_matrix(R_world_cam).wxyz
+
+
+def chase_pose(
+    subject: np.ndarray,
+    focus: np.ndarray,
+    *,
+    chase_distance: float = 15.0,
+    vertical_offset: float = 2.0,
+    up=(0.0, 0.0, 1.0),
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return ``(cam_pos, cam_wxyz, look_at)`` for a chase camera behind ``subject``.
+
+    The camera sits on the ray from ``focus`` through ``subject``, extended
+    ``chase_distance`` units past the subject, then lifted along world up by
+    ``vertical_offset``. It always looks at ``focus``.
+    """
+    subject = np.asarray(subject, dtype=np.float64)
+    focus = np.asarray(focus, dtype=np.float64)
+    up = np.asarray(up, dtype=np.float64)
+
+    ray = subject - focus
+    ray_norm = np.linalg.norm(ray)
+    if ray_norm < 1e-6:
+        cam_pos = subject + vertical_offset * up
+    else:
+        cam_pos = subject + chase_distance * (ray / ray_norm) + vertical_offset * up
+
+    wxyz = look_at_wxyz(cam_pos, focus, up)
+    return cam_pos, wxyz, focus
+
+
+def onboard_pose(
+    position: np.ndarray,
+    attitude_wxyz: np.ndarray,
+    R_sb: np.ndarray,
+    *,
+    forward_offset: float = 0.0,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return ``(cam_pos, cam_wxyz, look_at)`` for a sensor-mounted FPV camera.
+
+    Places the camera at ``position`` (shifted slightly forward along the sensor
+    boresight by ``forward_offset``) with orientation matching the sensor frame —
+    i.e. looking along the sensor boresight (+Z in sensor frame, mapped through
+    ``R_sb`` and the body attitude to world coordinates).
+    """
+    # R_body_to_world from the attitude quaternion (w, x, y, z)
+    w, x, y, z = attitude_wxyz
+    R_bw = np.array(
+        [
+            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
+            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
+            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
+        ],
+        dtype=np.float64,
+    )
+
+    # Sensor-to-world: columns are the sensor axes in world coords.
+    # R_sb is body-to-sensor, so R_sb.T is sensor-to-body.
+    R_sensor_to_world = R_bw @ R_sb.T
+
+    # Viser uses OpenCV convention (+X right, +Y DOWN, +Z forward).
+    # The sensor frame has +Y UP, so we rotate 180deg around the boresight (Z)
+    # to flip both X and Y, converting to the convention viser expects.
+    R_opencv_from_sensor = np.diag([-1.0, -1.0, 1.0])
+    R_cam_to_world = R_sensor_to_world @ R_opencv_from_sensor
+
+    wxyz = vtf.SO3.from_matrix(R_cam_to_world).wxyz
+    # Boresight is sensor +Z expressed in world frame (unchanged by the flip).
+    boresight_world = R_sensor_to_world[:, 2]
+    cam_pos = position + forward_offset * boresight_world
+    look_at = cam_pos + 10.0 * boresight_world
+    return cam_pos, wxyz, look_at
+
+
+def overview_pose(
+    positions: np.ndarray,
+    *,
+    azimuth: float = np.radians(135.0),
+    elevation: float = np.radians(25.0),
+    radius_margin: float = 0.75,
+    fov_deg: float = 60.0,
+    up=(0.0, 0.0, 1.0),
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return a static ``(cam_pos, cam_wxyz, look_at)`` that frames all ``positions``.
+
+    The camera is placed on a sphere around the centroid of ``positions``,
+    parameterized by ``azimuth`` (angle in XY from +X, CCW) and ``elevation``
+    (angle above the horizon). The radius is auto-computed so the full extent
+    fits within ``fov_deg``, then scaled by ``radius_margin``.
+    """
+    up = np.asarray(up, dtype=np.float64)
+    center = positions.mean(axis=0)
+    max_extent = np.max(np.linalg.norm(positions - center, axis=1))
+    half_fov_rad = np.radians(fov_deg / 2.0)
+    radius = max_extent / np.sin(half_fov_rad) * radius_margin
+
+    cos_el = np.cos(elevation)
+    cam_pos = center + radius * np.array(
+        [
+            cos_el * np.cos(azimuth),
+            cos_el * np.sin(azimuth),
+            np.sin(elevation),
+        ]
+    )
+    look_at = center.copy()
+    wxyz = look_at_wxyz(cam_pos, look_at, up)
+    return cam_pos, wxyz, look_at