Skip to content

Commit

Permalink
pydrake multibody: Add bindings for Toppra (RobotLocomotion#15836)
Browse files Browse the repository at this point in the history
  • Loading branch information
mpetersen94 authored Sep 30, 2021
1 parent 109fe5f commit dffb8c0
Show file tree
Hide file tree
Showing 3 changed files with 187 additions and 12 deletions.
2 changes: 1 addition & 1 deletion bindings/pydrake/multibody/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -293,12 +293,12 @@ drake_py_unittest(
drake_py_unittest(
name = "optimization_test",
timeout = "moderate",
tags = ["snopt"],
deps = [
":benchmarks_py",
":inverse_kinematics_py",
":optimization_py",
":parsing_py",
"//bindings/pydrake:trajectories_py",
"//bindings/pydrake/common/test_utilities",
"//bindings/pydrake/solvers:ipopt_py",
"//bindings/pydrake/solvers:snopt_py",
Expand Down
76 changes: 67 additions & 9 deletions bindings/pydrake/multibody/optimization_py.cc
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "drake/multibody/optimization/centroidal_momentum_constraint.h"
#include "drake/multibody/optimization/quaternion_integration_constraint.h"
#include "drake/multibody/optimization/static_equilibrium_problem.h"
#include "drake/multibody/optimization/toppra.h"

namespace drake {
namespace pydrake {
Expand All @@ -24,6 +25,19 @@ PYBIND11_MODULE(optimization, m) {
py::module::import("pydrake.multibody.plant");
py::module::import("pydrake.solvers.mathematicalprogram");

{
using Class = CalcGridPointsOptions;
constexpr auto& cls_doc = doc.CalcGridPointsOptions;
py::class_<Class>(m, "CalcGridPointsOptions", cls_doc.doc)
.def(ParamInit<Class>())
.def_readwrite("max_err", &Class::max_err, cls_doc.max_err.doc)
.def_readwrite("max_iter", &Class::max_iter, cls_doc.max_iter.doc)
.def_readwrite("max_seg_length", &Class::max_seg_length,
cls_doc.max_seg_length.doc)
.def_readwrite(
"min_points", &Class::min_points, cls_doc.min_points.doc);
}

{
using Class = CentroidalMomentumConstraint;
constexpr auto& cls_doc = doc.CentroidalMomentumConstraint;
Expand Down Expand Up @@ -59,6 +73,16 @@ PYBIND11_MODULE(optimization, m) {
AddValueInstantiation<ContactWrench>(m);
}

{
using Class = QuaternionEulerIntegrationConstraint;
constexpr auto& cls_doc = doc.QuaternionEulerIntegrationConstraint;
using Ptr = std::shared_ptr<Class>;
py::class_<Class, solvers::Constraint, Ptr>(
m, "QuaternionEulerIntegrationConstraint", cls_doc.doc)
.def(py::init<bool>(), py::arg("allow_quaternion_negation"),
cls_doc.ctor.doc);
}

{
using Class = StaticEquilibriumProblem;
constexpr auto& cls_doc = doc.StaticEquilibriumProblem;
Expand All @@ -84,15 +108,49 @@ PYBIND11_MODULE(optimization, m) {
}

{
using Class = QuaternionEulerIntegrationConstraint;
constexpr auto& cls_doc = doc.QuaternionEulerIntegrationConstraint;
using Ptr = std::shared_ptr<Class>;
py::class_<Class, solvers::Constraint, Ptr>(
m, "QuaternionEulerIntegrationConstraint", cls_doc.doc)
.def(py::init([](bool allow_quaternion_negation) {
return std::make_unique<Class>(allow_quaternion_negation);
}),
py::arg("allow_quaternion_negation"), cls_doc.ctor.doc);
using Class = ToppraDiscretization;
constexpr auto& cls_doc = doc.ToppraDiscretization;
py::enum_<Class>(m, "ToppraDiscretization", cls_doc.doc)
.value("kCollocation", Class::kCollocation, cls_doc.kCollocation.doc)
.value("kInterpolation", Class::kInterpolation,
cls_doc.kInterpolation.doc);
}

{
using Class = Toppra;
constexpr auto& cls_doc = doc.Toppra;
py::class_<Class>(m, "Toppra", cls_doc.doc)
.def(py::init<const Trajectory<double>&, const MultibodyPlant<double>&,
const Eigen::Ref<const Eigen::VectorXd>&>(),
py::arg("path"), py::arg("plant"), py::arg("gridpoints"),
cls_doc.ctor.doc)
.def_static("CalcGridPoints", &Class::CalcGridPoints, py::arg("path"),
py::arg("options"), cls_doc.CalcGridPoints.doc)
.def("SolvePathParameterization", &Class::SolvePathParameterization,
cls_doc.SolvePathParameterization.doc)
.def("AddJointVelocityLimit", &Class::AddJointVelocityLimit,
py::arg("lower_limit"), py::arg("upper_limit"),
cls_doc.AddJointVelocityLimit.doc)
.def("AddJointAccelerationLimit", &Class::AddJointAccelerationLimit,
py::arg("lower_limit"), py::arg("upper_limit"),
py::arg("discretization") = ToppraDiscretization::kInterpolation,
cls_doc.AddJointAccelerationLimit.doc)
.def("AddJointTorqueLimit", &Class::AddJointTorqueLimit,
py::arg("lower_limit"), py::arg("upper_limit"),
py::arg("discretization") = ToppraDiscretization::kInterpolation,
cls_doc.AddJointTorqueLimit.doc)
.def("AddFrameVelocityLimit", &Class::AddFrameVelocityLimit,
py::arg("constraint_frame"), py::arg("lower_limit"),
py::arg("upper_limit"), cls_doc.AddFrameVelocityLimit.doc)
.def("AddFrameTranslationalSpeedLimit",
&Class::AddFrameTranslationalSpeedLimit,
py::arg("constraint_frame"), py::arg("upper_limit"),
cls_doc.AddFrameTranslationalSpeedLimit.doc)
.def("AddFrameAccelerationLimit", &Class::AddFrameAccelerationLimit,
py::arg("constraint_frame"), py::arg("lower_limit"),
py::arg("upper_limit"),
py::arg("discretization") = ToppraDiscretization::kInterpolation,
cls_doc.AddFrameAccelerationLimit.doc);
}
}
} // namespace
Expand Down
121 changes: 119 additions & 2 deletions bindings/pydrake/multibody/test/optimization_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,21 @@
import numpy as np

from pydrake.multibody.optimization import (
CalcGridPointsOptions,
CentroidalMomentumConstraint,
QuaternionEulerIntegrationConstraint,
StaticEquilibriumProblem
StaticEquilibriumProblem,
Toppra,
ToppraDiscretization,
)
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import (
AddMultibodyPlantSceneGraph,
CoulombFriction
CoulombFriction,
MultibodyPlant,
)
from pydrake.multibody.tree import (
PrismaticJoint,
UnitInertia,
SpatialInertia,
)
Expand All @@ -27,6 +33,8 @@
)
from pydrake.math import RigidTransform
from pydrake.autodiffutils import AutoDiffXd
from pydrake.common import FindResourceOrThrow
from pydrake.trajectories import PiecewisePolynomial

Environment = namedtuple("Environment", [
"plant", "scene_graph", "diagram", "boxes", "ground_geometry_id",
Expand Down Expand Up @@ -107,6 +115,7 @@ def split_se3(q_se3):

class TestStaticEquilibriumProblem(unittest.TestCase):

@unittest.skipUnless(SnoptSolver().available(), "Requires Snopt")
def test_one_box(self):
# Test with a single box.
masses = [1.]
Expand Down Expand Up @@ -168,3 +177,111 @@ def test(self):
dut = QuaternionEulerIntegrationConstraint(
allow_quaternion_negation=True)
self.assertIsInstance(dut, mp.Constraint)


class TestToppra(unittest.TestCase):
def test_gridpoints(self):
path = PiecewisePolynomial.CubicWithContinuousSecondDerivatives(
[0., 1.], np.array([[0., -1.], [1, 1.1], [2, 3]]), np.zeros(3),
np.zeros(3))

options = CalcGridPointsOptions()
self.assertEqual(options.max_err, 1e-3)
options.max_err = 1e-2
self.assertEqual(options.max_iter, 100)
options.max_iter = 5
self.assertEqual(options.max_seg_length, 0.05)
options.max_seg_length = 0.1
self.assertEqual(options.min_points, 100)
options.min_points = 10

grid_points = Toppra.CalcGridPoints(path=path, options=options)
self.assertIsInstance(grid_points, np.ndarray)

def test_all_constraints(self):
plant = MultibodyPlant(0)
file_path = FindResourceOrThrow(
"drake/manipulation/models/iiwa_description/iiwa7/"
"iiwa7_no_collision.sdf")
iiwa_id = Parser(plant).AddModelFromFile(file_path, "iiwa")
plant.WeldFrames(plant.world_frame(),
plant.GetFrameByName("iiwa_link_0", iiwa_id))
plant.Finalize()

knots = np.array([[-1.57, 0.1, 0, -1.2, 0, 1.6, 0],
[-0.8, -0.6, 0.5, 0.1, -0.5, -0.1, -1]]).T
path = PiecewisePolynomial.CubicWithContinuousSecondDerivatives(
[0., 1.], knots, np.zeros(7), np.zeros(7))
gridpoints = Toppra.CalcGridPoints(path, CalcGridPointsOptions())
toppra = Toppra(path=path, plant=plant, gridpoints=gridpoints)

# Joint constraints
upper_limit = np.arange(7)
lower_limit = -upper_limit
frame = plant.GetFrameByName("iiwa_link_7")

constraint = toppra.AddJointVelocityLimit(lower_limit=lower_limit,
upper_limit=upper_limit)
self.assertIsInstance(constraint, mp.Binding[mp.BoundingBoxConstraint])
backward_con, forward_con = toppra.AddJointAccelerationLimit(
lower_limit=lower_limit, upper_limit=upper_limit,
discretization=ToppraDiscretization.kCollocation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])
backward_con, forward_con = toppra.AddJointAccelerationLimit(
lower_limit=lower_limit, upper_limit=upper_limit,
discretization=ToppraDiscretization.kInterpolation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])
backward_con, forward_con = toppra.AddJointTorqueLimit(
lower_limit=lower_limit, upper_limit=upper_limit,
discretization=ToppraDiscretization.kCollocation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])
backward_con, forward_con = toppra.AddJointTorqueLimit(
lower_limit=lower_limit, upper_limit=upper_limit,
discretization=ToppraDiscretization.kInterpolation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])

# Frame constraints
upper_limit = np.arange(6)
lower_limit = -upper_limit
frame = plant.GetFrameByName("iiwa_link_7")

constraint = toppra.AddFrameVelocityLimit(constraint_frame=frame,
lower_limit=lower_limit,
upper_limit=upper_limit)
self.assertIsInstance(constraint, mp.Binding[mp.BoundingBoxConstraint])
constraint = toppra.AddFrameTranslationalSpeedLimit(
constraint_frame=frame, upper_limit=1.)
self.assertIsInstance(constraint, mp.Binding[mp.BoundingBoxConstraint])
constraint = toppra.AddFrameAccelerationLimit(
constraint_frame=frame, lower_limit=lower_limit,
upper_limit=upper_limit,
discretization=ToppraDiscretization.kCollocation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])
constraint = toppra.AddFrameAccelerationLimit(
constraint_frame=frame, lower_limit=lower_limit,
upper_limit=upper_limit,
discretization=ToppraDiscretization.kInterpolation)
self.assertIsInstance(backward_con, mp.Binding[mp.LinearConstraint])
self.assertIsInstance(forward_con, mp.Binding[mp.LinearConstraint])

def test_solve(self):
plant = MultibodyPlant(0)
M_AAo_A = SpatialInertia(1, np.zeros(3), UnitInertia(1, 1, 1))
body = plant.AddRigidBody("body", M_AAo_A)
plant.AddJoint(PrismaticJoint(
"joint", plant.world_frame(), body.body_frame(), [1, 0, 0]))
plant.Finalize()

path = PiecewisePolynomial.CubicWithContinuousSecondDerivatives(
[0., 1.], np.array([[0., 1.]]), [0.], [0.])
gridpoints = Toppra.CalcGridPoints(path, CalcGridPointsOptions())
toppra = Toppra(path=path, plant=plant, gridpoints=gridpoints)

toppra.AddJointAccelerationLimit([-1.], [1.])
trajectory = toppra.SolvePathParameterization()
self.assertIsInstance(trajectory, PiecewisePolynomial)

0 comments on commit dffb8c0

Please sign in to comment.