Skip to content

Instantly share code, notes, and snippets.

@fferri

fferri/drake_example.cc

Created November 29, 2024 09:26
Show Gist options
  • Save fferri/fb26687f2138fbe572cb26f4317c8f6e to your computer and use it in GitHub Desktop.
Save fferri/fb26687f2138fbe572cb26f4317c8f6e to your computer and use it in GitHub Desktop.
Download ZIP
drake example
Raw
drake_example.cc
#include <string>
#include <vector>
#include <map>
#include <cmath>
#include <random>
#include <chrono>
#include <iostream>
#include <Eigen/Geometry>
#include <drake/geometry/drake_visualizer.h>
#include <drake/geometry/scene_graph.h>
#include <drake/math/rigid_transform.h>
#include <drake/multibody/parsing/parser.h>
#include <drake/multibody/plant/multibody_plant.h>
#include <drake/multibody/plant/multibody_plant_config_functions.h>
#include <drake/multibody/tree/prismatic_joint.h>
#include <drake/multibody/tree/revolute_joint.h>
#include <drake/multibody/tree/weld_joint.h>
#include <drake/systems/analysis/simulator.h>
#include <drake/systems/framework/diagram_builder.h>
#include <drake/systems/framework/leaf_system.h>
#include <drake/visualization/visualization_config_functions.h>
using std::pair;
using std::string;
using std::vector;
using std::map;
using std::optional;
using std::unique_ptr;
using std::shared_ptr;
using std::make_pair;
using std::make_unique;
using std::make_shared;
namespace e = Eigen;
namespace d = drake;
namespace drake {
namespace m = math;
namespace s = systems;
namespace g = geometry;
namespace mb = multibody;
namespace v = visualization;
}
namespace Eigen {
namespace rand {
template<typename T>
e::Vector3<T> vector(optional<std::uniform_real_distribution<T>> xd = {}, optional<std::uniform_real_distribution<T>> yd = {}, optional<std::uniform_real_distribution<T>> zd = {}) {
std::mt19937 _random_gen(std::random_device{}());
return {
xd.value_or(std::uniform_real_distribution<T>{-1, 1})(_random_gen),
yd.value_or(std::uniform_real_distribution<T>{-1, 1})(_random_gen),
zd.value_or(std::uniform_real_distribution<T>{-1, 1})(_random_gen)
};
}
template<typename T>
e::Quaternion<T> quaternion() {
std::mt19937 _random_gen(std::random_device{}());
std::uniform_real_distribution<T> dist{-1, 1};
return e::Quaternion<T>{
dist(_random_gen),
dist(_random_gen),
dist(_random_gen),
dist(_random_gen)
}.normalized();
}
}
}
template<typename T>
class Simulation {
public:
Simulation(const d::mb::MultibodyPlantConfig &cfg, shared_ptr<d::g::Meshcat> meshcat = {})
: plant_config_(cfg), psg_(d::mb::AddMultibodyPlant(cfg, &diagram_builder_)), meshcat_(meshcat) {}
void copyStateFrom(const Simulation &s) {
const d::mb::MultibodyPlant<T> &plant = s.psg_.plant;
const d::g::SceneGraph<T> &scene_graph = s.psg_.scene_graph;
const d::s::Context<T> &ctx = s.simulator_->get_context();
const d::s::Context<T> &pctx = plant.GetMyContextFromRoot(ctx);
#if 0
d::mb::ExportUrdf(plant, "tmp.urdf");
d::mb::Parser parser(&plant);
parser.AddModelFromFile("tmp.urdf");
#else
const d::g::SceneGraphInspector<T> &inspector = scene_graph.model_inspector();
std::map<const d::mb::RigidBody<T>*, const d::mb::RigidBody<T>*> bodyMap;
bodyMap[&plant.world_body()] = &psg_.plant.world_body();
for(int i = 1; i < plant.num_bodies(); i++) {
const d::mb::RigidBody<T> &body = plant.get_body(d::mb::BodyIndex(i));
d::mb::SpatialInertia<T> si = body.CalcSpatialInertiaInBodyFrame(pctx);
const d::mb::RigidBody<T> &newBody = psg_.plant.AddRigidBody(body.name(), si);
bodyMap[&body] = &newBody;
for(d::g::GeometryId id : plant.GetCollisionGeometriesForBody(body)) {
const d::g::Shape& geom = inspector.GetShape(id);
const string &n = inspector.GetName(id);
const d::g::ProximityProperties *pp = inspector.GetProximityProperties(id);
psg_.plant.RegisterCollisionGeometry(newBody, d::m::RigidTransformd(), geom, n, *pp);
}
for(d::g::GeometryId id : plant.GetVisualGeometriesForBody(body)) {
const d::g::Shape& geom = inspector.GetShape(id);
const string &n = inspector.GetName(id);
const d::g::IllustrationProperties *ip = inspector.GetIllustrationProperties(id);
psg_.plant.RegisterVisualGeometry(newBody, d::m::RigidTransformd(), geom, n, *ip);
}
setInitialBodyState(newBody, body.EvalPoseInWorld(pctx), body.EvalSpatialVelocityInWorld(pctx));
}
for(int i = 0; i < plant.num_joints(); i++) {
const d::mb::Joint<T> &joint = plant.get_joint(d::mb::JointIndex(i));
const string &name = joint.name();
const d::mb::RigidBody<T> &parent = *bodyMap[&joint.parent_body()];
d::m::RigidTransform<T> X_PF = joint.frame_on_parent().GetFixedPoseInBodyFrame();
const d::mb::RigidBody<T> &child = *bodyMap[&joint.child_body()];
d::m::RigidTransform<T> X_BM = joint.frame_on_child().GetFixedPoseInBodyFrame();
if(auto *j = dynamic_cast<const d::mb::RevoluteJoint<T>*>(&joint)) {
auto &newJoint = psg_.plant.template AddJoint<d::mb::RevoluteJoint>(name, parent, X_PF, child, X_BM, j->revolute_axis(), j->GetDamping(pctx));
setInitialJointState(newJoint, j->GetPositions(pctx)(0), j->GetVelocities(pctx)(0));
clearInitialBodyState(child);
} else if(auto *j = dynamic_cast<const d::mb::PrismaticJoint<T>*>(&joint)) {
auto &newJoint = psg_.plant.template AddJoint<d::mb::PrismaticJoint>(name, parent, X_PF, child, X_BM, j->translation_axis(), j->position_lower_limit(), j->position_upper_limit(), j->GetDamping(pctx));
setInitialJointState(newJoint, j->GetPositions(pctx)(0), j->GetVelocities(pctx)(0));
clearInitialBodyState(child);
} else if(auto *j = dynamic_cast<const d::mb::WeldJoint<T>*>(&joint)) {
auto &newJoint = psg_.plant.template AddJoint<d::mb::WeldJoint>(name, parent, X_PF, child, X_BM, j->X_FM());
clearInitialBodyState(child);
#if 0
} else {
const string &msg = fmt::format("unknown joint type at {} (name = {})", fmt::ptr(&joint), name);
std::cout << msg << std::endl;
throw std::runtime_error(msg);
#endif
}
}
#endif
}
void init() {
d::mb::MultibodyPlant<T> &plant = psg_.plant;
plant.Finalize();
if(meshcat_)
d::v::AddDefaultVisualization(&diagram_builder_, meshcat_);
diagram_ = diagram_builder_.Build();
simulator_ = make_unique<d::s::Simulator<T>>(*diagram_);
d::s::Context<T> &pctx = plant.GetMyMutableContextFromRoot(&simulator_->get_mutable_context());
for(const auto &p : initBodyState_) {
const auto &body = *p.first;
auto [pos, vel] = p.second;
plant.SetFreeBodyPose(&pctx, body, pos);
plant.SetFreeBodySpatialVelocity(&pctx, body, vel);
}
for(const auto &p : initJointState_) {
const auto &joint = *p.first;
auto [pos, vel] = p.second;
joint.SetPositions(&pctx, e::Vector<T, 1>{pos});
joint.SetVelocities(&pctx, e::Vector<T, 1>{vel});
}
simulator_->set_publish_every_time_step(true);
simulator_->set_target_realtime_rate(1.0);
simulator_->Initialize();
}
double getSimulationTime() {
return simulator_->get_context().get_time();
}
double getSimulationTimeStep() {
return plant_config_.time_step;
}
void step() {
double t = getSimulationTime();
double dt = getSimulationTimeStep();
simulator_->AdvanceTo(t + dt);
}
void step(const d::mb::JointActuator<T> &joint) {
if(getSimulationTime() > 1) {
d::mb::MultibodyPlant<T> &plant = psg_.plant;
d::s::Context<T> &pctx = plant.GetMyMutableContextFromRoot(&simulator_->get_mutable_context());
d::mb::MultibodyForces<T> f(plant);
joint.AddInOneForce(pctx, 0, 0.001, &f); //CRASH
}
step();
}
const d::mb::RigidBody<T>& world() const {
return psg_.plant.world_body();
}
enum class PrimitiveShapeType {
Cuboid,
Sphere,
Cylinder
};
const d::mb::RigidBody<T>& createPrimitiveShape(PrimitiveShapeType type, const string &name, const double &mass, const e::Vector3d &size, const e::Vector4d &color = {0.5, 0.5, 0.5, 1.0}) {
unique_ptr<d::g::Shape> shape;
d::mb::SpatialInertia<T> inertia = d::mb::SpatialInertia<T>::Zero();
switch(type) {
case PrimitiveShapeType::Cuboid:
shape = make_unique<d::g::Box>(size);
inertia = d::mb::SpatialInertia<T>::SolidBoxWithMass(mass, size(0), size(1), size(2));
break;
case PrimitiveShapeType::Sphere:
shape = make_unique<d::g::Sphere>(size(0));
inertia = d::mb::SpatialInertia<T>::SolidSphereWithMass(mass, size(0));
break;
case PrimitiveShapeType::Cylinder:
shape = make_unique<d::g::Cylinder>(size(0), size(2));
inertia = d::mb::SpatialInertia<T>::SolidCylinderWithMass(mass, size(0), size(2), e::Vector3<T>{0, 0, 1});
break;
default:
throw std::domain_error("invalid primitive shape type");
}
d::mb::MultibodyPlant<T> &plant = psg_.plant;
const d::mb::RigidBody<T> &body = plant.AddRigidBody(name, inertia);
d::mb::CoulombFriction<T> f(static_friction_, dynamic_friction_);
plant.RegisterCollisionGeometry(body, d::m::RigidTransformd(), *shape, name + "_collision", f);
plant.RegisterVisualGeometry(body, d::m::RigidTransformd(), *shape, name + "_visual", color);
return body;
}
const d::mb::RigidBody<T>& createMeshShape(const vector<e::Vector3d> &vertices, const vector<e::Vector3i> triangles, const string &name, const double &mass, const e::Vector4d &color = {0.5, 0.5, 0.5, 1.0}) {
std::ostringstream oss;
for(const auto& vertex : vertices)
oss << "v " << vertex(0) << " " << vertex(1) << " " << vertex(2) << "\n";
for(const auto& face : triangles)
oss << "f " << (face(0) + 1) << " " << (face(1) + 1) << " " << (face(2) + 1) << "\n";
d::MemoryFile mf{oss.str(), ".obj", "mesh"};
d::g::InMemoryMesh in_memory_mesh;
in_memory_mesh.mesh_file = mf;
unique_ptr<d::g::Shape> shape = make_unique<d::g::Mesh>(in_memory_mesh);
d::mb::SpatialInertia<T> inertia = d::mb::SpatialInertia<T>::MakeUnitary();
d::mb::MultibodyPlant<T> &plant = psg_.plant;
const d::mb::RigidBody<T>& body = plant.AddRigidBody(name, inertia);
d::mb::CoulombFriction<T> f(static_friction_, dynamic_friction_);
plant.RegisterCollisionGeometry(body, d::m::RigidTransformd(), *shape, name + "_collision", f);
plant.RegisterVisualGeometry(body, d::m::RigidTransformd(), *shape, name + "_visual", color);
return body;
}
enum class JointType {
Revolute,
Prismatic,
Fixed
};
const d::mb::Joint<T>& createJoint(JointType type, const string &name, const d::mb::RigidBody<T> &parent, const d::mb::RigidBody<T> &child, d::m::RigidTransform<T> X_PF = d::m::RigidTransform<T>::Identity(), d::m::RigidTransform<T> X_BM = d::m::RigidTransform<T>::Identity(), d::m::RigidTransform<T> X_FM = d::m::RigidTransform<T>::Identity()) {
d::mb::MultibodyPlant<T> &plant = psg_.plant;
// frames' names:
// P = parent, B = child
// joint: F (attaches to parent), M (attaches to child)
clearInitialBodyState(child);
switch(type) {
case JointType::Revolute:
return plant.template AddJoint<d::mb::RevoluteJoint>(name, parent, X_PF, child, X_BM, e::Vector3<T>{0, 0, 1}, 0);
case JointType::Prismatic:
return plant.template AddJoint<d::mb::PrismaticJoint>(name, parent, X_PF, child, X_BM, e::Vector3<T>{0, 0, 1}, 0, 1, 0);
case JointType::Fixed:
return plant.template AddJoint<d::mb::WeldJoint>(name, parent, X_PF, child, X_BM, X_FM);
default:
throw std::domain_error("invalid joint type");
}
}
const d::mb::JointActuator<T>& createJointActuator(const string &name, const d::mb::Joint<T> &joint, double effort_limit = std::numeric_limits<T>::infinity()) {
d::mb::MultibodyPlant<T> &plant = psg_.plant;
return plant.AddJointActuator(name, joint, effort_limit);
}
void setInitialBodyState(const d::mb::RigidBody<T> &body, const d::m::RigidTransform<T> &pose, const d::mb::SpatialVelocity<T> &vel) {
initBodyState_[&body] = make_pair(pose, vel);
}
void clearInitialBodyState(const d::mb::RigidBody<T> &body) {
initBodyState_.erase(&body);
}
void setInitialJointState(const d::mb::Joint<T> &joint, T pos, T vel) {
initJointState_[&joint] = make_pair(pos, vel);
}
void clearInitialJointState(const d::mb::Joint<T> &joint) {
initJointState_.erase(&joint);
}
static void test() {
shared_ptr<d::g::Meshcat> meshcat = make_shared<d::g::Meshcat>();
d::mb::MultibodyPlantConfig cfg;
cfg.time_step = 0.005; // s
cfg.stiction_tolerance = 1.0E-4; // m/s
cfg.penetration_allowance = 1.0E-4;
cfg.discrete_contact_approximation = "sap"; // tamsi || sap || similar || lagged
unique_ptr<Simulation> ps = make_unique<Simulation>(cfg, meshcat);
const d::mb::RigidBody<T> &ground = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "ground", 1.0, {10, 10, 0.1}, {0.3, 0.3, 0.8, 1});
ps->createJoint(Simulation::JointType::Fixed, "ground_fixed", ps->world(), ground);
// double pendulum:
const d::mb::RigidBody<T> &c1 = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "c1", 1.0, {0.1, 0.1, 0.1}, {0.3, 0.3, 0.3, 1});
ps->createJoint(Simulation::JointType::Fixed, "f1", ps->world(), c1, {}, {}, d::m::RigidTransform<T>{e::Vector3<T>{0, 0, 1}});
const d::mb::RigidBody<T> &c2 = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "c2", 1.0, {0.1, 0.1, 0.5}, {0.3, 0.3, 0.3, 1});
auto qzx = e::Quaternion<T>::FromTwoVectors(e::Vector3<T>{1, 0, 0}, e::Vector3<T>{0, 0, 1});
const d::mb::Joint<T> &j1 = ps->createJoint(Simulation::JointType::Revolute, "j1", c1, c2, d::m::RigidTransform<T>{qzx, e::Vector3<T>{0.05, 0, 0}}, d::m::RigidTransform<T>{qzx, e::Vector3<T>{-0.05, 0, 0.25-0.05}});
ps->setInitialJointState(j1, M_PI / 2., 0.1);
const d::mb::RigidBody<T> &c3 = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "c3", 1.0, {0.1, 0.1, 0.5}, {0.3, 0.3, 0.3, 1});
const d::mb::Joint<T> &j2 = ps->createJoint(Simulation::JointType::Revolute, "j2", c2, c3, d::m::RigidTransform<T>{qzx, e::Vector3<T>{0.05, 0, -0.25+0.05}}, d::m::RigidTransform<T>{qzx, e::Vector3<T>{-0.05, 0, 0.25-0.05}});
// rotating lever:
const d::mb::RigidBody<T> &c4 = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cylinder, "c4", 1.0, {0.1, 0.1, 0.1}, {0.3, 0.3, 0.3, 1});
ps->createJoint(Simulation::JointType::Fixed, "f2", ps->world(), c4, {}, {}, d::m::RigidTransform<T>{e::Vector3<T>{0.5, 0, 0.2}});
//ps->setInitialBodyState(c4, d::m::RigidTransform<T>{e::Vector3<T>{0.5, 0, 0.05}}, {});
const d::mb::RigidBody<T> &c5 = ps->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "c5", 1.0, {0.1, 1.0, 0.1}, {0.3, 0.3, 0.3, 1});
const d::mb::Joint<T> &j3 = ps->createJoint(Simulation::JointType::Revolute, "j3", c4, c5, d::m::RigidTransform<T>{e::Vector3<T>{0.0, 0.0, 0.05}}, d::m::RigidTransform<T>{e::Vector3<T>{0.0, 0.5-0.05, -0.05}});
const d::mb::JointActuator<T> &a3 = ps->createJointActuator("a3", j3);
ps->init();
while(true) {
double t = ps->getSimulationTime();
double dt = ps->getSimulationTimeStep();
// spawn a shape every second:
if(t > 0 && fmod(t, 1) < dt) {
unique_ptr<Simulation> ps1 = make_unique<Simulation>(cfg, meshcat);
ps1->copyStateFrom(*ps);
static int i = 0;
static double sqrt2 = std::sqrt(2), sqrt3 = std::sqrt(3), sqrt6 = std::sqrt(6), k = 0.1;
static std::vector<e::Vector3d> vrt{
{k * std::sqrt(8. / 9.), 0, k * -1. / 3.},
{k * -std::sqrt(2. / 9.), k * std::sqrt(2. / 3.), k * -1. / 3.},
{k * -std::sqrt(2. / 9.), k * -std::sqrt(2. / 3.), k * -1. / 3.},
{0, 0, k * 1},
};
static std::vector<e::Vector3i> idx{{0, 1, 2}, {0, 2, 3}, {0, 3, 1}, {1, 3, 2}};
const d::mb::RigidBody<T> &b = (++i % 2)
? ps1->createPrimitiveShape(Simulation::PrimitiveShapeType::Cuboid, "cube-" + std::to_string(i), 1.0, {0.1, 0.1, 0.1}, {0.8, 0.3, 0.3, 1})
: ps1->createMeshShape(vrt, idx, "mesh-" + std::to_string(i), 1.0, {0.8, 0.3, 0.3, 1});
ps1->setInitialBodyState(b,
d::m::RigidTransform<T>(e::rand::quaternion<T>(), e::Vector3d{0, 0, 1.3}),
d::mb::SpatialVelocity<T>(4 * e::rand::vector<T>(), 0.2 * e::rand::vector<T>() + e::Vector3d{0, 0, 2})
);
ps1->init();
ps = std::move(ps1);
}
ps->step(a3);
}
}
private:
d::s::DiagramBuilder<T> diagram_builder_;
const d::mb::MultibodyPlantConfig &plant_config_;
d::mb::AddMultibodyPlantSceneGraphResult<T> psg_;
unique_ptr<d::s::Diagram<T>> diagram_;
unique_ptr<d::s::Simulator<T>> simulator_;
shared_ptr<d::g::Meshcat> meshcat_;
map<const d::mb::RigidBody<T>*, pair<d::m::RigidTransform<T>, d::mb::SpatialVelocity<T>>> initBodyState_;
map<const d::mb::Joint<T>*, pair<T, T>> initJointState_;
const double static_friction_ = 0.8;
const double dynamic_friction_ = 0.5; // (not used in stepped simulations)
};
int main() {
Simulation<double>::test();
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment