Drake KinematicTrajectoryOptimization PathConstraint Autodiff Benchmark

benchmark.py

"""Time the real PathConstraint from KinematicTrajectoryOptimization.

Obtains the actual internal PathConstraint (anonymous namespace, not directly
constructible) via the public AddPathPositionConstraint() API, then times its
double-scalar Eval().  The measured call is identical before and after the
commit -- a plain float64 array in, a plain float64 array out -- so the pydrake
marshalling cost is the same in both builds and cancels in the comparison.
"""
import os, sys, time
import numpy as np
from pydrake.geometry import Box, Role
from pydrake.math import RigidTransform
from pydrake.multibody.inverse_kinematics import MinimumDistanceLowerBoundConstraint
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph, CoulombFriction
from pydrake.planning import KinematicTrajectoryOptimization
from pydrake.solvers import LinearConstraint
from pydrake.systems.framework import DiagramBuilder

IIWA = "package://drake_models/iiwa_description/urdf/iiwa14_spheres_collision.urdf"
OBSTACLES = [(0.6,0,0.4),(0,0.6,0.5),(-0.5,0,0.6),(0,-0.6,0.3)]
NUM_SAMPLES, NUM_REPS = 200, 20

def build_plant():
    b = DiagramBuilder()
    plant, sg = AddMultibodyPlantSceneGraph(b, time_step=0.0)
    parser = Parser(plant)
    # drake_models is a lazily-fetched external repo; the runfiles tree only
    # symlinks what the borrowed test target needed, so register it explicitly.
    models_dir = os.environ.get("DRAKE_MODELS_DIR")
    if models_dir:
        pm = parser.package_map()
        if pm.Contains("drake_models"):
            pm.Remove("drake_models")
        pm.Add("drake_models", models_dir)
    model = parser.AddModelsFromUrl(IIWA)[0]
    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base"))
    for i, xyz in enumerate(OBSTACLES):
        body = plant.AddRigidBody(f"obstacle{i}", model)
        plant.RegisterCollisionGeometry(body, RigidTransform(), Box(.2,.2,.2),
                                        f"obstacle{i}", CoulombFriction(1,1))
        plant.WeldFrames(plant.world_frame(), body.body_frame(),
                         RigidTransform(np.array(xyz)))
    plant.Finalize()
    d = b.Build(); ctx = d.CreateDefaultContext()
    return plant, sg, plant.GetMyContextFromRoot(ctx)

def time_eval(evaluator, samples, reps=NUM_REPS):
    for x in samples[:10]:
        evaluator.Eval(x)
    t0 = time.perf_counter()
    for _ in range(reps):
        for x in samples:
            evaluator.Eval(x)
    return (time.perf_counter() - t0) / (reps * len(samples)) * 1e6

def main():
    plant, sg, plant_context = build_plant()
    nq = plant.num_positions()
    insp = sg.model_inspector()
    ngeo = insp.NumGeometriesWithRole(Role.kProximity)
    ncand = len(insp.GetCollisionCandidates())
    rng = np.random.default_rng(0)

    # Baseline: pydrake Eval() marshalling cost for a same-sized double vector.
    trivial_order = 4
    lc = LinearConstraint(np.ones((1, trivial_order*nq)), [-np.inf], [np.inf])
    ov = time_eval(lc, [rng.uniform(-1,1,trivial_order*nq) for _ in range(NUM_SAMPLES)])

    print(f"scene: nq={nq}, {ngeo} collision geometries "
          f"({ngeo-len(OBSTACLES)} robot spheres + {len(OBSTACLES)} obstacle boxes), "
          f"{ncand} candidate pairs")
    print(f"pydrake Eval() marshalling baseline (double, same size): {ov:.2f} us")
    print(f"{'influence_off':>13} {'num_vars':>9} {'active':>7} {'us/eval':>9} {'net':>9}")

    qport = plant.get_geometry_query_input_port()
    for off in (0.01, 0.1, 0.5):
        trajopt = KinematicTrajectoryOptimization(nq, 10)
        con = MinimumDistanceLowerBoundConstraint(
            plant, 0.01, plant_context, influence_distance_offset=off)
        binding = trajopt.AddPathPositionConstraint(con, 0.5)
        path_constraint = binding.evaluator()   # <-- the real PathConstraint
        nvars = path_constraint.num_vars()
        samples = [rng.uniform(-1.0, 1.0, nvars) for _ in range(NUM_SAMPLES)]

        # How many collision pairs are active at the paths these x's produce.
        counts = []
        for x in samples:
            path_constraint.Eval(x)     # leaves plant_context at the path point
            counts.append(len(qport.Eval(plant_context)
                              .ComputeSignedDistancePairwiseClosestPoints(0.01+off)))
        active = float(np.mean(counts))

        t = time_eval(path_constraint, samples)
        print(f"{off:>13} {nvars:>9} {active:>7.1f} {t:>9.2f} {t-ov:>9.2f}")

if __name__ == "__main__":
    main()