Skip to content

Commit

Permalink
Add python binding for static equlibrium problem. (RobotLocomotion#14140
Browse files Browse the repository at this point in the history
)

Also use ComputeSignedDistancePairwiseClosestPoints() to avoid redundant computation.
  • Loading branch information
hongkai-dai authored Oct 9, 2020
1 parent 6c470a2 commit 33cbce4
Show file tree
Hide file tree
Showing 9 changed files with 353 additions and 115 deletions.
30 changes: 30 additions & 0 deletions bindings/pydrake/multibody/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,22 @@ drake_pybind_library(
],
)

drake_pybind_library(
name = "optimization_py",
cc_deps = [
"//bindings/pydrake:documentation_pybind",
"//bindings/pydrake/common:default_scalars_pybind",
"//bindings/pydrake/common:value_pybind",
],
cc_srcs = ["optimization_py.cc"],
package_info = PACKAGE_INFO,
py_deps = [
":module_py",
":plant_py",
"//bindings/pydrake/solvers:mathematicalprogram_py",
],
)

drake_jupyter_py_binary(
name = "examples/door_hinge_inspector",
add_test_rule = 1,
Expand All @@ -161,6 +177,7 @@ PY_LIBRARIES_WITH_INSTALL = [
":benchmarks_py",
":inverse_kinematics_py",
":math_py",
":optimization_py",
":parsing_py",
":plant_py",
":tree_py",
Expand Down Expand Up @@ -270,6 +287,19 @@ drake_py_unittest(
],
)

drake_py_unittest(
name = "optimization_test",
timeout = "moderate",
deps = [
":benchmarks_py",
":inverse_kinematics_py",
":optimization_py",
":parsing_py",
"//bindings/pydrake/common/test_utilities",
"//bindings/pydrake/solvers:ipopt_py",
],
)

add_pybind_coverage_data()

add_lint_tests_pydrake()
1 change: 1 addition & 0 deletions bindings/pydrake/multibody/all.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
# Normal symbols.
from .inverse_kinematics import * # noqa
from .math import * # noqa
from .optimization import * # noqa
from .parsing import * # noqa
from .plant import * # noqa
from .tree import * # noqa
Expand Down
64 changes: 64 additions & 0 deletions bindings/pydrake/multibody/optimization_py.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
#include "pybind11/eigen.h"
#include "pybind11/pybind11.h"

#include "drake/bindings/pydrake/common/default_scalars_pybind.h"
#include "drake/bindings/pydrake/common/value_pybind.h"
#include "drake/bindings/pydrake/documentation_pybind.h"
#include "drake/bindings/pydrake/pydrake_pybind.h"
#include "drake/multibody/optimization/static_equilibrium_problem.h"

namespace drake {
namespace pydrake {

namespace {
PYBIND11_MODULE(optimization, m) {
// NOLINTNEXTLINE(build/namespaces): Emulate placement in namespace.
using namespace drake::multibody;
constexpr auto& doc = pydrake_doc.drake.multibody;

m.doc() = "Optimization module for MultibodyPlant motion planning";

py::module::import("pydrake.math");
py::module::import("pydrake.multibody.plant");
py::module::import("pydrake.solvers.mathematicalprogram");

{
py::class_<ContactWrench>(m, "ContactWrench", doc.ContactWrench.doc)
.def_readonly("bodyA_index", &ContactWrench::bodyA_index,
doc.ContactWrench.bodyA_index.doc)
.def_readonly("bodyB_index", &ContactWrench::bodyB_index,
doc.ContactWrench.bodyB_index.doc)
.def_readonly(
"p_WCb_W", &ContactWrench::p_WCb_W, doc.ContactWrench.p_WCb_W.doc)
.def_readonly(
"F_Cb_W", &ContactWrench::F_Cb_W, doc.ContactWrench.F_Cb_W.doc);
AddValueInstantiation<ContactWrench>(m);
}

{
using Class = StaticEquilibriumProblem;
constexpr auto& cls_doc = doc.StaticEquilibriumProblem;
py::class_<Class>(m, "StaticEquilibriumProblem", cls_doc.doc)
.def(py::init<const MultibodyPlant<AutoDiffXd>*,
systems::Context<AutoDiffXd>*,
const std::set<
std::pair<geometry::GeometryId, geometry::GeometryId>>&>(),
py::arg("plant"), py::arg("context"),
py::arg("ignored_collision_pairs"),
// Keep alive, reference: `self` keeps `plant` and `context` alive.
py::keep_alive<1, 2>(), py::keep_alive<1, 3>(), cls_doc.ctor.doc)
.def("get_mutable_prog", &Class::get_mutable_prog,
py_rvp::reference_internal, cls_doc.get_mutable_prog.doc)
.def("prog", &Class::prog, py_rvp::reference_internal, cls_doc.prog.doc)
.def("q_vars", &Class::q_vars, cls_doc.q_vars.doc)
.def("u_vars", &Class::u_vars, cls_doc.u_vars.doc)
.def("GetContactWrenchSolution", &Class::GetContactWrenchSolution,
py::arg("result"), cls_doc.GetContactWrenchSolution.doc)
.def("UpdateComplementarityTolerance",
&Class::UpdateComplementarityTolerance, py::arg("tol"),
cls_doc.UpdateComplementarityTolerance.doc);
}
}
} // namespace
} // namespace pydrake
} // namespace drake
176 changes: 176 additions & 0 deletions bindings/pydrake/multibody/test/optimization_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,176 @@
import unittest
import typing

import numpy as np

from pydrake.multibody.optimization import StaticEquilibriumProblem
from pydrake.multibody.plant import (
AddMultibodyPlantSceneGraph,
CoulombFriction
)
from pydrake.multibody.tree import (
UnitInertia,
SpatialInertia,
)
import pydrake.multibody.inverse_kinematics as ik
import pydrake.solvers.mathematicalprogram as mp
from pydrake.solvers.ipopt import IpoptSolver
from pydrake.systems.framework import DiagramBuilder_
from pydrake.geometry import (
Box,
Sphere,
)
from pydrake.math import RigidTransform
from pydrake.autodiffutils import AutoDiffXd


def construct_environment(masses: typing.List, box_sizes: typing.List):
"""
Construct an environment with many free boxes.
@param masses masses[i] is the mass of box i.
@param box_sizes box_sizes[i] is the size of box i.
"""
builder = DiagramBuilder_[AutoDiffXd]()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.0)
# Add the ground as a big box.
ground_box = plant.AddRigidBody(
"ground", SpatialInertia(1, np.array([0, 0, 0]), UnitInertia(1, 1, 1)))
X_WG = RigidTransform([0, 0, -0.05])
ground_geometry_id = plant.RegisterCollisionGeometry(
ground_box, RigidTransform(), Box(10, 10, 0.1), "ground",
CoulombFriction(0.9, 0.8))
plant.RegisterVisualGeometry(
ground_box, RigidTransform(), Box(10, 10, 0.1), "ground",
[0.5, 0.5, 0.5, 1.])
plant.WeldFrames(plant.world_frame(), ground_box.body_frame(), X_WG)

# Add boxes
assert isinstance(masses, list)
assert isinstance(box_sizes, list)
assert len(masses) == len(box_sizes)
num_boxes = len(masses)
boxes = []
boxes_geometry_id = []
for i in range(num_boxes):
box_name = f"box{i}"
boxes.append(plant.AddRigidBody(
box_name, SpatialInertia(
masses[i], np.array([0, 0, 0]), UnitInertia(1, 1, 1))))
box_shape = Box(box_sizes[i][0], box_sizes[i][1], box_sizes[i][2])
boxes_geometry_id.append(plant.RegisterCollisionGeometry(
boxes[i], RigidTransform(), box_shape, f"{box_name}_box",
CoulombFriction(0.9, 0.8)))
plant.RegisterVisualGeometry(
ground_box, RigidTransform(), box_shape, f"{box_name}_box",
[0.5, 0.5, 0.5, 1.])
# Add small spheres at the corners of the box.
vertices = np.array([
[1, 1, 1], [1, 1, -1], [1, -1, 1], [1, -1, -1], [-1, 1, 1],
[-1, -1, 1], [-1, 1, -1], [-1, -1, -1]]) *\
np.tile(box_sizes[i]/2, (8, 1))
sphere_shape = Sphere(0.001)
for j in range(8):
plant.RegisterCollisionGeometry(
boxes[i], RigidTransform(vertices[j]), sphere_shape,
f"{box_name}_sphere{j}", CoulombFriction(0.9, 0.8))
plant.RegisterVisualGeometry(
boxes[i], RigidTransform(vertices[j]), sphere_shape,
f"{box_name}_sphere{j}", [0.5, 0.5, 0.5, 1])

plant.Finalize()
diagram = builder.Build()

return plant, scene_graph, diagram, boxes, ground_geometry_id,\
boxes_geometry_id


def split_se3(q_se3):
"""
q_se3 is the quaternion/position of a floating body.
"""
assert q_se3.shape == (7,)
quat = q_se3[:4]
xyz = q_se3[4:]
return quat, xyz


class TestStaticEquilibriumProblem(unittest.TestCase):

def test_one_box(self):
# Test with a single box.
masses = [1.]
box_sizes = [np.array([0.1, 0.1, 0.1])]
plant, scene_graph, diagram, boxes, ground_geometry_id,\
boxes_geometry_id = construct_environment(
masses, box_sizes)
diagram_context = diagram.CreateDefaultContext()
plant_context = diagram.GetMutableSubsystemContext(
plant, diagram_context)
# Ignore the collision between box-box, since currently Drake doesn't
# support box-box collision with AutoDiffXd.
ignored_collision_pairs = {(
ground_geometry_id, boxes_geometry_id[0])}
dut = StaticEquilibriumProblem(
plant, plant_context, ignored_collision_pairs)

# Add the constraint that the quaternion should have unit length.
box_quat, box_xyz = split_se3(dut.q_vars())
ik.AddUnitQuaternionConstraintOnPlant(
plant, dut.q_vars(), dut.get_mutable_prog())

dut.get_mutable_prog().SetInitialGuess(
box_quat, np.array([0.5, 0.5, 0.5, 0.5]))
dut.get_mutable_prog().SetInitialGuess(
box_xyz, np.array([0, 0, 0.3]))

# Now set the complementarity tolerance.
dut.UpdateComplementarityTolerance(0.002)
result = mp.Solve(dut.prog())
self.assertTrue(result.is_success())
q_sol = result.GetSolution(dut.q_vars())
box_quat_sol, box_xyz_sol = split_se3(q_sol)
self.assertAlmostEqual(np.linalg.norm(box_quat_sol), 1.)
np.testing.assert_allclose(box_xyz_sol[2], 0.05, atol=0.002)

contact_wrenches = dut.GetContactWrenchSolution(result)
self.assertEqual(len(contact_wrenches), 8)

def test_two_boxes(self):
# test with 2 boxes.
masses = [1., 2]
box_sizes = [np.array([0.1, 0.1, 0.1]), np.array([0.1, 0.1, 0.2])]
plant, scene_graph, diagram, boxes, ground_geometry_id,\
boxes_geometry_id = construct_environment(
masses, box_sizes)
diagram_context = diagram.CreateDefaultContext()
plant_context = diagram.GetMutableSubsystemContext(
plant, diagram_context)
# Ignore the collision between box-box, since currently Drake doesn't
# support box-box collision with AutoDiffXd.
ignored_collision_pairs = set()
for i in range(len(boxes_geometry_id)):
ignored_collision_pairs.add(tuple((
ground_geometry_id, boxes_geometry_id[i])))
for j in range(i+1, len(boxes_geometry_id)):
ignored_collision_pairs.add(tuple(
(boxes_geometry_id[i], boxes_geometry_id[j])))

dut = StaticEquilibriumProblem(
plant, plant_context, ignored_collision_pairs)

# Add the constraint that the quaternion should have unit length.
ik.AddUnitQuaternionConstraintOnPlant(
plant, dut.q_vars(), dut.get_mutable_prog())
for i in range(len(boxes)):
box_quaternion_start = 7 * i
box_quat, box_xyz = split_se3(
dut.q_vars()[box_quaternion_start:box_quaternion_start+7])
dut.get_mutable_prog().SetInitialGuess(
box_quat, np.array([0.5, 0.5, 0.5, 0.5]))
dut.get_mutable_prog().SetInitialGuess(
box_xyz, np.array([0, 0, 0.3 + 0.2 * i]))

dut.UpdateComplementarityTolerance(0.001)
ipopt_solver = IpoptSolver()
result = ipopt_solver.Solve(dut.prog())
self.assertTrue(result.is_success())
16 changes: 10 additions & 6 deletions multibody/optimization/contact_wrench.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,22 +14,26 @@ namespace multibody {
*/
struct ContactWrench {
/**
* @param bodyA_index_in The index of Body A.
* @param bodyB_index_in The index of Body B.
* @param p_WCb_W_in The position of the point Cb (where the wrench is
* applied) expressed in the world frame W.
* @param F_Cb_W_in The wrench (spatial force) applied at point Cb from Body A
* to Body B, measured in the world frame.
* Refer to the documentation for each attribute.
*/
ContactWrench(BodyIndex bodyA_index_in, BodyIndex bodyB_index_in,
Eigen::Vector3d p_WCb_W_in, SpatialForce<double> F_Cb_W_in)
: bodyA_index(bodyA_index_in),
bodyB_index(bodyB_index_in),
p_WCb_W(std::move(p_WCb_W_in)),
F_Cb_W(std::move(F_Cb_W_in)) {}

/** The index of Body A.*/
BodyIndex bodyA_index;
/** The index of Body B.*/
BodyIndex bodyB_index;
/** The position of the point Cb (where the wrench is applied) expressed in
* the world frame W.
*/
Eigen::Vector3d p_WCb_W;
/** F_Cb_W_in The wrench (spatial force) applied at point Cb from Body A
* to Body B, measured in the world frame.
*/
SpatialForce<double> F_Cb_W;
};
} // namespace multibody
Expand Down
29 changes: 10 additions & 19 deletions multibody/optimization/static_equilibrium_constraint.cc
Original file line number Diff line number Diff line change
Expand Up @@ -62,14 +62,15 @@ void StaticEquilibriumConstraint::DoEval(
}
const auto& query_object =
query_port.Eval<geometry::QueryObject<AutoDiffXd>>(*context_);
const std::vector<geometry::SignedDistancePair<AutoDiffXd>>
signed_distance_pairs =
query_object.ComputeSignedDistancePairwiseClosestPoints();
const geometry::SceneGraphInspector<AutoDiffXd>& inspector =
query_object.inspector();
const int lambda_start_index_in_x =
plant_->num_positions() + plant_->num_actuated_dofs();
for (const auto& signed_distance_pair : signed_distance_pairs) {
for (const auto& [contact_pair, evaluator] :
contact_pair_to_wrench_evaluator_) {
const geometry::SignedDistancePair<AutoDiffXd> signed_distance_pair =
query_object.ComputeSignedDistancePairClosestPoints(
contact_pair.first(), contact_pair.second());
const geometry::FrameId frame_A_id =
inspector.GetFrameId(signed_distance_pair.id_A);
const geometry::FrameId frame_B_id =
Expand Down Expand Up @@ -101,29 +102,19 @@ void StaticEquilibriumConstraint::DoEval(
frameB, p_BCb, plant_->world_frame(),
plant_->world_frame(), &Jv_V_WCb);

const SortedPair<geometry::GeometryId> contact_pair(
signed_distance_pair.id_A, signed_distance_pair.id_B);
// Find the lambda corresponding to the geometry pair (id_A, id_B).
const auto it = contact_pair_to_wrench_evaluator_.find(contact_pair);
if (it == contact_pair_to_wrench_evaluator_.end()) {
throw std::runtime_error(
"The input argument contact_pair_to_wrench_evaluator in the "
"StaticEquilibriumConstraint constructor doesn't include all "
"possible contact pairs.");
}

VectorX<AutoDiffXd> lambda(
it->second.contact_wrench_evaluator->num_lambda());
evaluator.contact_wrench_evaluator->num_lambda());

for (int i = 0; i < lambda.rows(); ++i) {
lambda(i) = x(lambda_start_index_in_x +
it->second.lambda_indices_in_all_lambda[i]);
evaluator.lambda_indices_in_all_lambda[i]);
}

AutoDiffVecXd F_AB_W;
it->second.contact_wrench_evaluator->Eval(
it->second.contact_wrench_evaluator->ComposeVariableValues(*context_,
lambda),
evaluator.contact_wrench_evaluator->Eval(
evaluator.contact_wrench_evaluator->ComposeVariableValues(*context_,
lambda),
&F_AB_W);

// By definition, F_AB_W is the contact wrench applied to id_B from id_A,
Expand Down
Loading

0 comments on commit 33cbce4

Please sign in to comment.