import numpy as np
import os
import pydrake
from pydrake.geometry import (
MeshcatVisualizerCpp,
MeshcatVisualizerParams,
Role,
StartMeshcat,
)
from pydrake.math import RigidTransform, RollPitchYaw
from pydrake.multibody.meshcat import JointSliders
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import AddMultibodyPlantSceneGraph
from pydrake.systems.analysis import Simulator
from pydrake.systems.framework import DiagramBuilder, LeafSystem, PortDataType, BasicVector
from pydrake.systems.primitives import ConstantVectorSource
from pydrake.systems.controllers import PidController
# Custom Data Bus System: (Courtesy of Eric C.)
class custom_leaf_system(LeafSystem):
def __init__(self, num_in, num_out):
LeafSystem.__init__(self)
input_port = self.DeclareInputPort(
"input",
PortDataType.kVectorValued,
num_in,
)
# The output is re-arranging the input:
def output(context, output):
x = input_port.Eval(context)
y = np.array([x[1], x[0]], dtype=float)
output.set_value(y)
output_port = self.DeclareVectorOutputPort(
"output",
BasicVector(num_out),
output,
)
meshcat = StartMeshcat()
def model_inspector(filename):
meshcat.Delete()
meshcat.DeleteAddedControls()
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)
# Load the file into the plant/scene_graph.
parser = Parser(plant)
parser.AddModelFromFile(filename)
plant.Finalize()
# Add two visualizers, one to publish the "visual" geometry, and one to
# publish the "collision" geometry.
visual = MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat,
MeshcatVisualizerParams(role=Role.kPerception, prefix="visual"))
collision = MeshcatVisualizerCpp.AddToBuilder(
builder, scene_graph, meshcat,
MeshcatVisualizerParams(role=Role.kProximity, prefix="collision"))
# Disable the collision geometry at the start; it can be enabled by the
# checkbox in the meshcat controls.
meshcat.SetProperty("collision", "visible", False)
# Set the timeout to a small number in test mode. Otherwise, JointSliders
# will run until "Stop JointSliders" button is clicked.
default_interactive_timeout = 1.0 if "TEST_SRCDIR" in os.environ else None
sliders = builder.AddSystem(JointSliders(meshcat, plant))
diagram = builder.Build()
sliders.Run(diagram, default_interactive_timeout)
# Press the 'Stop JointSliders' button in MeshCat to continue.
model_location = "/work/cart_pole.sdf"
# Test to see if model works:
model_inspector(model_location)
def create_scene(filename, sim_time_step=0.0001):
# Clean up MeshCat.
meshcat.Delete()
meshcat.DeleteAddedControls()
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(
builder, time_step=sim_time_step)
parser = Parser(plant)
# Load model:
model_location = filename
parser.AddModelFromFile(model_location)
# Add Gravity:
plant.mutable_gravity_field().set_gravity_vector([0, 0, -9.8])
plant.Finalize()
# Set Initial Position of Joints:
eps_ = 0.3
joint_0 = plant.GetJointByName('prismatic_joint')
joint_0.set_default_translation(0.0)
joint_1 = plant.GetJointByName('revolute_joint')
joint_1.set_default_angle(np.pi - eps_)
context = plant.CreateDefaultContext()
# Add visualizer to visualize the geometries.
visualizer = MeshcatVisualizerCpp.AddToBuilder(
builder, scene_graph, meshcat,
MeshcatVisualizerParams(role=Role.kPerception, prefix="visual"))
# Make Diagram:
# Construct and Add Controller:
kp = np.array([0, 5], dtype=float)
kd = np.array([0, 1], dtype=float)
ki = np.array([0, 10], dtype=float)
pid_controller = builder.AddSystem(PidController(kp, ki, kd))
# Constant Desired State Trajectory Block:
desired_trajectory = builder.AddSystem(ConstantVectorSource(np.array([0, np.pi, 0, 0], dtype=float)))
# Construct and Add Custom Data Bus:
data_bus = custom_leaf_system(2, 2)
builder.AddSystem(data_bus)
# Connect Desired Trajectory and PID:
builder.Connect(
desired_trajectory.get_output_port(),
pid_controller.get_input_port_desired_state()
)
# Connect Plants States to PID Estimated States:
builder.Connect(
plant.get_state_output_port(),
pid_controller.get_input_port_estimated_state()
)
# Connect PID Controller to Data Bus:
builder.Connect(
pid_controller.get_output_port_control(),
data_bus.get_input_port()
)
# Connect Data Bus Output to Plant Input Actuation:
builder.Connect(
data_bus.get_output_port(),
plant.get_actuation_input_port()
)
diagram = builder.Build()
return diagram, visualizer, plant, context
def initialize_simulation(diagram):
simulator = Simulator(diagram)
simulator.Initialize()
simulator.set_target_realtime_rate(1.0)
return simulator
def run_simulation(model_location, sim_time_step):
diagram, visualizer, plant, context = create_scene(model_location, sim_time_step)
simulator = initialize_simulation(diagram)
visualizer.StartRecording()
simulator.AdvanceTo(5.0)
visualizer.PublishRecording()
# Run the simulation:
model_location = "/work/cart_pole.sdf"
run_simulation(model_location, sim_time_step=0.0001)