forked from pz4kybsvg/Conception
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
330 lines
13 KiB
330 lines
13 KiB
#include "drake/planning/trajectory_optimization/direct_collocation.h"
|
|
|
|
#include <cstddef>
|
|
#include <stdexcept>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "drake/math/autodiff.h"
|
|
#include "drake/math/autodiff_gradient.h"
|
|
|
|
namespace drake {
|
|
namespace planning {
|
|
namespace trajectory_optimization {
|
|
|
|
using math::AreAutoDiffVecXdEqual;
|
|
using solvers::Binding;
|
|
using solvers::Constraint;
|
|
using solvers::MathematicalProgram;
|
|
using solvers::VectorXDecisionVariable;
|
|
using systems::Context;
|
|
using systems::FixedInputPortValue;
|
|
using systems::InputPort;
|
|
using systems::InputPortIndex;
|
|
using systems::InputPortSelection;
|
|
using systems::PortDataType;
|
|
using systems::System;
|
|
using trajectories::PiecewisePolynomial;
|
|
|
|
namespace {
|
|
int CheckAndReturnStates(int states) {
|
|
if (states <= 0) {
|
|
throw std::logic_error(
|
|
"This system doesn't have any continuous states. DirectCollocation "
|
|
"only makes sense for systems with continuous-time dynamics.");
|
|
}
|
|
return states;
|
|
}
|
|
|
|
typedef std::pair<std::unique_ptr<System<AutoDiffXd>>,
|
|
std::unique_ptr<Context<AutoDiffXd>>>
|
|
OwnedPair;
|
|
|
|
OwnedPair MakeAutoDiffXd(const systems::System<double>& system,
|
|
const systems::Context<double>& context) {
|
|
auto system_ad = System<double>::ToAutoDiffXd(system);
|
|
auto context_ad = system_ad->CreateDefaultContext();
|
|
// TODO(russt): Add support for time-varying dynamics OR check for
|
|
// time-invariance.
|
|
context_ad->SetTimeStateAndParametersFrom(context);
|
|
system_ad->FixInputPortsFrom(system, context, context_ad.get());
|
|
return OwnedPair{std::move(system_ad), std::move(context_ad)};
|
|
}
|
|
|
|
} // namespace
|
|
|
|
DirectCollocationConstraint::DirectCollocationConstraint(
|
|
const System<double>& system, const Context<double>& context,
|
|
std::variant<InputPortSelection, InputPortIndex> input_port_index,
|
|
bool assume_non_continuous_states_are_fixed)
|
|
: DirectCollocationConstraint(
|
|
MakeAutoDiffXd(system, context),
|
|
nullptr, // System
|
|
nullptr, nullptr, nullptr, // Contexts
|
|
context.num_continuous_states(),
|
|
system.get_input_port_selection(input_port_index)
|
|
? system.get_input_port_selection(input_port_index)->size()
|
|
: 0,
|
|
input_port_index, assume_non_continuous_states_are_fixed) {}
|
|
|
|
DirectCollocationConstraint::DirectCollocationConstraint(
|
|
const systems::System<AutoDiffXd>& system,
|
|
systems::Context<AutoDiffXd>* context_sample,
|
|
systems::Context<AutoDiffXd>* context_next_sample,
|
|
systems::Context<AutoDiffXd>* context_collocation,
|
|
std::variant<systems::InputPortSelection, systems::InputPortIndex>
|
|
input_port_index,
|
|
bool assume_non_continuous_states_are_fixed)
|
|
: DirectCollocationConstraint(
|
|
OwnedPair{nullptr, nullptr}, &system, context_sample,
|
|
context_next_sample, context_collocation,
|
|
context_sample->num_continuous_states(),
|
|
system.get_input_port_selection(input_port_index)
|
|
? system.get_input_port_selection(input_port_index)->size()
|
|
: 0,
|
|
input_port_index, assume_non_continuous_states_are_fixed) {}
|
|
|
|
DirectCollocationConstraint::DirectCollocationConstraint(
|
|
OwnedPair owned_pair, const systems::System<AutoDiffXd>* system,
|
|
systems::Context<AutoDiffXd>* context_sample,
|
|
systems::Context<AutoDiffXd>* context_next_sample,
|
|
systems::Context<AutoDiffXd>* context_collocation, int num_states,
|
|
int num_inputs,
|
|
std::variant<InputPortSelection, InputPortIndex> input_port_index,
|
|
bool assume_non_continuous_states_are_fixed)
|
|
: Constraint(CheckAndReturnStates(num_states),
|
|
1 + (2 * num_states) + (2 * num_inputs),
|
|
Eigen::VectorXd::Zero(num_states),
|
|
Eigen::VectorXd::Zero(num_states)),
|
|
owned_system_(std::move(owned_pair.first)),
|
|
owned_context_(std::move(owned_pair.second)),
|
|
system_(owned_system_ ? *owned_system_ : *system),
|
|
context_sample_(owned_context_ ? owned_context_.get() : context_sample),
|
|
context_next_sample_(owned_context_ ? owned_context_.get()
|
|
: context_next_sample),
|
|
context_collocation_(owned_context_ ? owned_context_.get()
|
|
: context_collocation),
|
|
input_port_(system_.get_input_port_selection(input_port_index)),
|
|
num_states_(num_states),
|
|
num_inputs_(num_inputs) {
|
|
system_.ValidateContext(context_sample_);
|
|
system_.ValidateContext(context_next_sample_);
|
|
system_.ValidateContext(context_collocation_);
|
|
if (!assume_non_continuous_states_are_fixed) {
|
|
DRAKE_THROW_UNLESS(context_sample_->has_only_continuous_state());
|
|
}
|
|
|
|
if (input_port_) {
|
|
// Verify that the input port is not abstract valued.
|
|
if (input_port_->get_data_type() == PortDataType::kAbstractValued) {
|
|
throw std::logic_error(
|
|
"The specified input port is abstract-valued, and this constraint "
|
|
"only supports vector-valued input ports. Did you perhaps forget to "
|
|
"pass a non-default `input_port_index` argument?");
|
|
}
|
|
}
|
|
}
|
|
|
|
void DirectCollocationConstraint::dynamics(const AutoDiffVecXd& state,
|
|
const AutoDiffVecXd& input,
|
|
Context<AutoDiffXd>* context,
|
|
AutoDiffVecXd* xdot) const {
|
|
if (input_port_ &&
|
|
(!input_port_->HasValue(*context) ||
|
|
!AreAutoDiffVecXdEqual(input, input_port_->Eval(*context)))) {
|
|
input_port_->FixValue(context, input);
|
|
}
|
|
if (!AreAutoDiffVecXdEqual(
|
|
state, context->get_continuous_state_vector().CopyToVector())) {
|
|
context->SetContinuousState(state);
|
|
}
|
|
*xdot = system_.EvalTimeDerivatives(*context).CopyToVector();
|
|
}
|
|
|
|
void DirectCollocationConstraint::DoEval(
|
|
const Eigen::Ref<const Eigen::VectorXd>& x, Eigen::VectorXd* y) const {
|
|
AutoDiffVecXd y_t;
|
|
Eval(x.cast<AutoDiffXd>(), &y_t);
|
|
*y = math::ExtractValue(y_t);
|
|
}
|
|
|
|
// The format of the input to the eval() function is the
|
|
// tuple { timestep, state 0, state 1, input 0, input 1 },
|
|
// which has a total length of 1 + 2*num_states + 2*num_inputs.
|
|
void DirectCollocationConstraint::DoEval(
|
|
const Eigen::Ref<const AutoDiffVecXd>& x, AutoDiffVecXd* y) const {
|
|
DRAKE_ASSERT(x.size() == 1 + (2 * num_states_) + (2 * num_inputs_));
|
|
|
|
// Extract our input variables:
|
|
// h - current time (breakpoint)
|
|
// x0, x1 state vector at time steps k, k+1
|
|
// u0, u1 input vector at time steps k, k+1
|
|
const AutoDiffXd h = x(0);
|
|
const auto x0 = x.segment(1, num_states_);
|
|
const auto x1 = x.segment(1 + num_states_, num_states_);
|
|
const auto u0 = x.segment(1 + (2 * num_states_), num_inputs_);
|
|
const auto u1 = x.segment(1 + (2 * num_states_) + num_inputs_, num_inputs_);
|
|
|
|
AutoDiffVecXd xdot0;
|
|
dynamics(x0, u0, context_sample_, &xdot0);
|
|
|
|
AutoDiffVecXd xdot1;
|
|
dynamics(x1, u1, context_next_sample_, &xdot1);
|
|
|
|
// Cubic interpolation to get xcol and xdotcol.
|
|
const AutoDiffVecXd xcol = 0.5 * (x0 + x1) + h / 8 * (xdot0 - xdot1);
|
|
const AutoDiffVecXd xdotcol = -1.5 * (x0 - x1) / h - .25 * (xdot0 + xdot1);
|
|
|
|
AutoDiffVecXd g;
|
|
dynamics(xcol, 0.5 * (u0 + u1), context_collocation_, &g);
|
|
*y = xdotcol - g;
|
|
}
|
|
|
|
void DirectCollocationConstraint::DoEval(
|
|
const Eigen::Ref<const VectorX<symbolic::Variable>>&,
|
|
VectorX<symbolic::Expression>*) const {
|
|
throw std::logic_error(
|
|
"DirectCollocationConstraint does not support symbolic evaluation.");
|
|
}
|
|
|
|
Binding<Constraint> AddDirectCollocationConstraint(
|
|
std::shared_ptr<DirectCollocationConstraint> constraint,
|
|
const Eigen::Ref<const VectorXDecisionVariable>& timestep,
|
|
const Eigen::Ref<const VectorXDecisionVariable>& state,
|
|
const Eigen::Ref<const VectorXDecisionVariable>& next_state,
|
|
const Eigen::Ref<const VectorXDecisionVariable>& input,
|
|
const Eigen::Ref<const VectorXDecisionVariable>& next_input,
|
|
MathematicalProgram* prog) {
|
|
DRAKE_DEMAND(timestep.size() == 1);
|
|
DRAKE_DEMAND(state.size() == constraint->num_states());
|
|
DRAKE_DEMAND(next_state.size() == constraint->num_states());
|
|
DRAKE_DEMAND(input.size() == constraint->num_inputs());
|
|
DRAKE_DEMAND(next_input.size() == constraint->num_inputs());
|
|
return prog->AddConstraint(constraint,
|
|
{timestep, state, next_state, input, next_input});
|
|
}
|
|
|
|
DirectCollocation::DirectCollocation(
|
|
const System<double>* system, const Context<double>& context,
|
|
int num_time_samples, double minimum_timestep, double maximum_timestep,
|
|
std::variant<InputPortSelection, InputPortIndex> input_port_index,
|
|
bool assume_non_continuous_states_are_fixed,
|
|
solvers::MathematicalProgram* prog)
|
|
: MultipleShooting(
|
|
system->get_input_port_selection(input_port_index)
|
|
? system->get_input_port_selection(input_port_index)->size()
|
|
: 0,
|
|
CheckAndReturnStates(context.num_continuous_states()),
|
|
num_time_samples, minimum_timestep, maximum_timestep, prog),
|
|
system_(system),
|
|
context_(context.Clone()),
|
|
input_port_index_(input_port_index),
|
|
sample_contexts_(num_time_samples) {
|
|
system->ValidateContext(context);
|
|
if (!assume_non_continuous_states_are_fixed) {
|
|
DRAKE_DEMAND(context.has_only_continuous_state());
|
|
}
|
|
|
|
auto system_and_context = MakeAutoDiffXd(*system, context);
|
|
system_ad_ = std::move(system_and_context.first);
|
|
context_ad_ = std::move(system_and_context.second);
|
|
|
|
// Allocated contexts for each sample time. We share contexts across multiple
|
|
// constraints in order to exploit caching (the dynamics at time k are
|
|
// evaluated both in constraint k and k+1). Note that the constraints cannot
|
|
// be evaluated in parallel.
|
|
for (int i = 0; i < N(); ++i) {
|
|
sample_contexts_[i] = context_ad_->Clone();
|
|
}
|
|
// We don't expect to have cache hits for the collocation contexts, so can
|
|
// just use the one context_ad.
|
|
|
|
// Add the dynamic constraints.
|
|
// For N-1 timesteps, add a constraint which depends on the breakpoint
|
|
// along with the state and input vectors at that breakpoint and the
|
|
// next.
|
|
for (int i = 0; i < N() - 1; ++i) {
|
|
auto constraint = std::make_shared<DirectCollocationConstraint>(
|
|
*system_ad_, sample_contexts_[i].get(), sample_contexts_[i + 1].get(),
|
|
context_ad_.get(), input_port_index,
|
|
assume_non_continuous_states_are_fixed);
|
|
this->prog()
|
|
.AddConstraint(constraint,
|
|
{h_vars().segment<1>(i),
|
|
x_vars().segment(i * num_states(), num_states() * 2),
|
|
u_vars().segment(i * num_inputs(), num_inputs() * 2)})
|
|
.evaluator()
|
|
->set_description(
|
|
fmt::format("collocation constraint for segment {}", i));
|
|
}
|
|
}
|
|
|
|
void DirectCollocation::DoAddRunningCost(const symbolic::Expression& g) {
|
|
// Trapezoidal integration:
|
|
// sum_{i=0...N-2} h_i/2.0 * (g_i + g_{i+1}), or
|
|
// g_0*h_0/2.0 + [sum_{i=1...N-2} g_i*(h_{i-1} + h_i)/2.0] +
|
|
// g_{N-1}*h_{N-2}/2.0.
|
|
|
|
prog().AddCost(SubstitutePlaceholderVariables(g * h_vars()(0) / 2, 0));
|
|
for (int i = 1; i <= N() - 2; i++) {
|
|
prog().AddCost(SubstitutePlaceholderVariables(
|
|
g * (h_vars()(i - 1) + h_vars()(i)) / 2, i));
|
|
}
|
|
prog().AddCost(SubstitutePlaceholderVariables(
|
|
g * h_vars()(N() - 2) / 2, N() - 1));
|
|
}
|
|
|
|
PiecewisePolynomial<double> DirectCollocation::ReconstructInputTrajectory(
|
|
const solvers::MathematicalProgramResult& result) const {
|
|
const InputPort<double>* input_port =
|
|
system_->get_input_port_selection(input_port_index_);
|
|
if (!input_port) {
|
|
return PiecewisePolynomial<double>();
|
|
}
|
|
|
|
Eigen::VectorXd times = GetSampleTimes(result);
|
|
std::vector<double> times_vec(N());
|
|
std::vector<Eigen::MatrixXd> inputs(N());
|
|
|
|
for (int i = 0; i < N(); i++) {
|
|
times_vec[i] = times(i);
|
|
inputs[i] = result.GetSolution(input(i));
|
|
}
|
|
return PiecewisePolynomial<double>::FirstOrderHold(times_vec, inputs);
|
|
}
|
|
|
|
PiecewisePolynomial<double> DirectCollocation::ReconstructStateTrajectory(
|
|
const solvers::MathematicalProgramResult& result) const {
|
|
Eigen::VectorXd times = GetSampleTimes(result);
|
|
std::vector<double> times_vec(N());
|
|
std::vector<Eigen::MatrixXd> states(N());
|
|
std::vector<Eigen::MatrixXd> derivatives(N());
|
|
|
|
// Provide a fixed value for the input port and keep an alias around.
|
|
const InputPort<double>* input_port =
|
|
system_->get_input_port_selection(input_port_index_);
|
|
FixedInputPortValue* input_port_value{nullptr};
|
|
if (input_port) {
|
|
input_port_value = &input_port->FixValue(
|
|
context_.get(),
|
|
system_->AllocateInputVector(*input_port)->get_value());
|
|
}
|
|
|
|
for (int i = 0; i < N(); i++) {
|
|
times_vec[i] = times(i);
|
|
states[i] = result.GetSolution(state(i));
|
|
if (input_port) {
|
|
input_port_value->GetMutableVectorData<double>()->SetFromVector(
|
|
result.GetSolution(input(i)));
|
|
}
|
|
context_->SetContinuousState(states[i]);
|
|
derivatives[i] = system_->EvalTimeDerivatives(*context_).CopyToVector();
|
|
}
|
|
return PiecewisePolynomial<double>::CubicHermite(times_vec, states,
|
|
derivatives);
|
|
}
|
|
|
|
} // namespace trajectory_optimization
|
|
} // namespace planning
|
|
} // namespace drake
|