def num_positions_velocities_example():
plant = MultibodyPlant(time_step=0.0)
Parser(plant).AddModelFromFile(
FindResourceOrThrow(
"drake/examples/manipulation_station/models/061_foam_brick.sdf"))
plant.Finalize()
context = plant.CreateDefaultContext()
print(context)
print(f"plant.num_positions() = {plant.num_positions()}")
print(f"plant.num_velocities() = {plant.num_velocities()}")
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("sdf_path", help="path to sdf")
parser.add_argument("--interactive", action='store_true')
MeshcatVisualizer.add_argparse_argument(parser)
args = parser.parse_args()
builder = DiagramBuilder()
scene_graph = builder.AddSystem(SceneGraph())
plant = MultibodyPlant(time_step=0.01)
plant.RegisterAsSourceForSceneGraph(scene_graph)
parser = Parser(plant)
model = parser.AddModelFromFile(args.sdf_path)
plant.Finalize()
if args.meshcat:
meshcat = ConnectMeshcatVisualizer(
builder, output_port=scene_graph.get_query_output_port(),
zmq_url=args.meshcat, open_browser=args.open_browser)
if args.interactive:
# Add sliders to set positions of the joints.
sliders = builder.AddSystem(JointSliders(robot=plant))
else:
q0 = plant.GetPositions(plant.CreateDefaultContext())
sliders = builder.AddSystem(ConstantVectorSource(q0))
to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(plant))
builder.Connect(sliders.get_output_port(0), to_pose.get_input_port())
builder.Connect(
to_pose.get_output_port(),
scene_graph.get_source_pose_port(plant.get_source_id()))
diagram = builder.Build()
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
simulator.AdvanceTo(1E6)
import matplotlib.pyplot as plt
from pydrake.all import (DiagramBuilder, DirectCollocation, MultibodyPlant,
MultibodyPositionToGeometryPose, Parser,
PiecewisePolynomial, PlanarSceneGraphVisualizer,
SceneGraph, Simulator, Solve, TrajectorySource)
from underactuated import FindResource
plant = MultibodyPlant(time_step=0.0)
scene_graph = SceneGraph()
plant.RegisterAsSourceForSceneGraph(scene_graph)
file_name = FindResource("models/cartpole.urdf")
Parser(plant).AddModelFromFile(file_name)
plant.Finalize()
context = plant.CreateDefaultContext()
dircol = DirectCollocation(
plant,
context,
num_time_samples=21,
minimum_timestep=0.1,
maximum_timestep=0.4,
input_port_index=plant.get_actuation_input_port().get_index())
dircol.AddEqualTimeIntervalsConstraints()
initial_state = (0., 0., 0., 0.)
dircol.AddBoundingBoxConstraint(initial_state, initial_state,
dircol.initial_state())
# More elegant version is blocked by drake #8315:
# dircol.AddLinearConstraint(dircol.initial_state() == initial_state)
class ProprioceptionSystem(pydrake.systems.framework.LeafSystem):
"""
Calculates J, M, Cv, etc. based on arm state.
"""
def __init__(self):
pydrake.systems.framework.LeafSystem.__init__(self)
# Physical parameters
self.manipulator_plant = MultibodyPlant(config.DT)
manipulator.data["add_plant_function"](self.manipulator_plant)
self.manipulator_plant.Finalize()
self.manipulator_plant_context = \
self.manipulator_plant.CreateDefaultContext()
self.nq_manipulator = \
self.manipulator_plant.get_actuation_input_port().size()
self.DeclareVectorInputPort("state",
BasicVector(self.nq_manipulator * 2))
self.DeclareVectorOutputPort(
"q", pydrake.systems.framework.BasicVector(self.nq_manipulator),
self.output_q)
self.DeclareVectorOutputPort(
"v", pydrake.systems.framework.BasicVector(self.nq_manipulator),
self.output_v)
self.DeclareVectorOutputPort(
"tau_g",
pydrake.systems.framework.BasicVector(self.nq_manipulator),
self.output_tau_g)
self.DeclareVectorOutputPort(
"M",
pydrake.systems.framework.BasicVector(
self.nq_manipulator * self.nq_manipulator), self.output_M)
self.DeclareVectorOutputPort(
"Cv", pydrake.systems.framework.BasicVector(self.nq_manipulator),
self.output_Cv)
self.DeclareVectorOutputPort(
"J",
pydrake.systems.framework.BasicVector(6 * self.nq_manipulator),
self.output_J)
self.DeclareVectorOutputPort(
"J_translational",
pydrake.systems.framework.BasicVector(3 * self.nq_manipulator),
self.output_J_translational)
self.DeclareVectorOutputPort(
"J_rotational",
pydrake.systems.framework.BasicVector(3 * self.nq_manipulator),
self.output_J_rotational)
self.DeclareVectorOutputPort("Jdot_qdot",
pydrake.systems.framework.BasicVector(6),
self.output_Jdot_qdot)
# =========================== CALC FUNCTIONS ==========================
def calc_q(self, context):
state = self.GetInputPort("state").Eval(context)
q = state[:self.nq_manipulator]
return q
def calc_v(self, context):
state = self.GetInputPort("state").Eval(context)
v = state[self.nq_manipulator:]
return v
def calc_J(self, context):
self.update_plant(context)
contact_body = self.manipulator_plant.GetBodyByName(
manipulator.data["contact_body_name"])
J = self.manipulator_plant.CalcJacobianSpatialVelocity(
self.manipulator_plant_context, JacobianWrtVariable.kV,
contact_body.body_frame(), [0, 0, 0],
self.manipulator_plant.world_frame(),
self.manipulator_plant.world_frame())
return J
# ========================== OUTPUT FUNCTIONS =========================
def output_q(self, context, output):
q = self.calc_q(context)
output.SetFromVector(q.flatten())
def output_v(self, context, output):
v = self.calc_v(context)
output.SetFromVector(v.flatten())
def output_tau_g(self, context, output):
self.update_plant(context)
tau_g = self.manipulator_plant.CalcGravityGeneralizedForces(
self.manipulator_plant_context)
output.SetFromVector(tau_g.flatten())
def output_M(self, context, output):
self.update_plant(context)
M = self.manipulator_plant.CalcMassMatrixViaInverseDynamics(
self.manipulator_plant_context)
output.SetFromVector(M.flatten())
def output_Cv(self, context, output):
self.update_plant(context)
Cv = self.manipulator_plant.CalcBiasTerm(
self.manipulator_plant_context)
output.SetFromVector(Cv.flatten())
def output_J(self, context, output):
J = self.calc_J(context)
output.SetFromVector(J.flatten())
def output_J_translational(self, context, output):
J = self.calc_J(context)
J_translational = J[3:, :]
output.SetFromVector(J_translational.flatten())
def output_J_rotational(self, context, output):
J = self.calc_J(context)
J_rotational = J[:3, :]
output.SetFromVector(J_rotational.flatten())
def output_Jdot_qdot(self, context, output):
self.update_plant(context)
contact_body = self.manipulator_plant.GetBodyByName(
manipulator.data["contact_body_name"])
Jdot_qdot_raw = self.manipulator_plant.CalcBiasSpatialAcceleration(
self.manipulator_plant_context, JacobianWrtVariable.kV,
contact_body.body_frame(), [0, 0, 0],
self.manipulator_plant.world_frame(),
self.manipulator_plant.world_frame())
Jdot_qdot = list(Jdot_qdot_raw.rotational()) + \
list(Jdot_qdot_raw.translational())
output.SetFromVector(Jdot_qdot)
# ========================== OTHER FUNCTIONS ==========================
def update_plant(self, context):
self.manipulator_plant.SetPositions(self.manipulator_plant_context,
self.calc_q(context))
self.manipulator_plant.SetVelocities(self.manipulator_plant_context,
self.calc_v(context))
class Visualizer():
def __init__(self, urdf):
self._create_plant(urdf)
def _create_plant(self, urdf):
self.plant = MultibodyPlant(time_step=0.0)
self.scenegraph = SceneGraph()
self.plant.RegisterAsSourceForSceneGraph(self.scenegraph)
self.model_index = Parser(self.plant).AddModelFromFile(FindResource(urdf))
self.builder = DiagramBuilder()
self.builder.AddSystem(self.scenegraph)
def _finalize_plant(self):
self.plant.Finalize()
def _add_trajectory_source(self, traj):
# Create the trajectory source
source = self.builder.AddSystem(TrajectorySource(traj))
pos2pose = self.builder.AddSystem(MultibodyPositionToGeometryPose(self.plant, input_multibody_state=True))
# Wire the source to the scene graph
self.builder.Connect(source.get_output_port(0), pos2pose.get_input_port())
self.builder.Connect(pos2pose.get_output_port(), self.scenegraph.get_source_pose_port(self.plant.get_source_id()))
def _add_renderer(self):
renderer_name = "renderer"
self.scenegraph.AddRenderer(renderer_name, MakeRenderEngineVtk(RenderEngineVtkParams()))
# Add a camera
depth_camera = DepthRenderCamera(
RenderCameraCore(
renderer_name,
CameraInfo(width=640, height=480, fov_y=np.pi/4),
ClippingRange(0.01, 10.0),
RigidTransform()),
DepthRange(0.01,10.0)
)
world_id = self.plant.GetBodyFrameIdOrThrow(self.plant.world_body().index())
X_WB = xyz_rpy_deg([4,0,0],[-90,0,90])
sensor = RgbdSensor(world_id, X_PB=X_WB, depth_camera=depth_camera)
self.builder.AddSystem(sensor)
self.builder.Connect(self.scenegraph.get_query_output_port(), sensor.query_object_input_port())
def _connect_visualizer(self):
DrakeVisualizer.AddToBuilder(self.builder, self.scenegraph)
self.meshcat = ConnectMeshcatVisualizer(self.builder, self.scenegraph, zmq_url="new", open_browser=False)
self.meshcat.vis.jupyter_cell()
def _make_visualization(self, stop_time):
simulator = Simulator(self.builder.Build())
simulator.Initialize()
self.meshcat.vis.render_static()
# Set simulator context
simulator.get_mutable_context().SetTime(0.0)
# Start recording and simulate the diagram
self.meshcat.reset_recording()
self.meshcat.start_recording()
simulator.AdvanceTo(stop_time)
# Publish the recording
self.meshcat.publish_recording()
# Render
self.meshcat.vis.render_static()
input("View visualization. Press <ENTER> to end")
print("Finished")
def visualize_trajectory(self, xtraj=None):
self._finalize_plant()
print("Adding trajectory source")
xtraj = self._check_trajectory(xtraj)
self._add_trajectory_source(xtraj)
print("Adding renderer")
self._add_renderer()
print("Connecting to MeshCat")
self._connect_visualizer()
print("Making visualization")
self._make_visualization(xtraj.end_time())
def _check_trajectory(self, traj):
if traj is None:
plant_context = self.plant.CreateDefaultContext()
pose = self.plant.GetPositions(plant_context)
pose = np.column_stack((pose, pose))
pose = zero_pad_rows(pose, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold([0., 1.], pose)
elif type(traj) is np.ndarray:
if traj.ndim == 1:
traj = np.column_stack(traj, traj)
traj = zero_pad_rows(traj, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold([0.,1.], traj)
elif type(traj) is PiecewisePolynomial:
breaks = traj.get_segment_times()
values = traj.vector_values(breaks)
values = zero_pad_rows(values, self.plant.num_positions() + self.plant.num_velocities())
return PiecewisePolynomial.FirstOrderHold(breaks, values)
else:
raise ValueError("Trajectory must be a piecewise polynomial, an ndarray, or None")
class iiwa_sys():
def __init__(self,
builder,
dt=5e-4,
N=150,
params=None,
trj_decay=0.7,
x_w_cov=1e-5,
door_angle_ref=1.0,
visualize=False):
self.plant_derivs = MultibodyPlant(time_step=dt)
parser = Parser(self.plant_derivs)
self.derivs_iiwa, _, _ = self.add_models(self.plant_derivs,
parser,
params=params)
self.plant_derivs.Finalize()
self.plant_derivs_context = self.plant_derivs.CreateDefaultContext()
self.plant_derivs.get_actuation_input_port().FixValue(
self.plant_derivs_context, [0., 0., 0., 0., 0., 0., 0.])
null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
self.plant_derivs.GetInputPort("applied_generalized_force").FixValue(
self.plant_derivs_context, null_force)
self.plant, scene_graph = AddMultibodyPlantSceneGraph(builder,
time_step=dt)
parser = Parser(self.plant, scene_graph)
self.iiwa, self.hinge, self.bushing = self.add_models(self.plant,
parser,
params=params)
self.plant.Finalize(
) # Finalize will assign ports for compatibility w/ the scene_graph; could be cause of the issue w/ first order taylor.
self.meshcat = ConnectMeshcatVisualizer(builder,
scene_graph,
zmq_url=zmq_url)
self.sym_derivs = False # If system should use symbolic derivatives; if false, autodiff
self.custom_sim = False # If rollouts should be gathered with sys.dyn() calls
nq = self.plant.num_positions()
nv = self.plant.num_velocities()
self.n_x = nq + nv
self.n_u = self.plant.num_actuators()
self.n_y = self.plant.get_state_output_port(self.iiwa).size()
self.N = N
self.dt = dt
self.decay = trj_decay
self.V = 1e-2 * np.ones(self.n_y)
self.W = np.concatenate((1e-7 * np.ones(nq), 1e-4 * np.ones(nv)))
self.W0 = np.concatenate((1e-9 * np.ones(nq), 1e-6 * np.ones(nv)))
self.x_w_cov = x_w_cov
self.door_angle_ref = door_angle_ref
self.q0 = np.array(
[-3.12, -0.17, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
#self.q0 = np.array([-3.12, -0.27, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
self.x0 = np.concatenate((self.q0, np.zeros(nv)))
self.door_index = None
self.phi = {}
def ports_init(self, context):
self.plant_context = self.plant.GetMyMutableContextFromRoot(context)
self.plant.SetPositionsAndVelocities(self.plant_context, self.x0)
#door_angle = self.plant.GetPositionsFromArray(self.hinge, self.x0[:8])
self.door_index = self.plant.GetJointByName(
'right_door_hinge').position_start()
null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
self.plant.GetInputPort("applied_generalized_force").FixValue(
self.plant_context, null_force)
self.plant.get_actuation_input_port().FixValue(
self.plant_context, [0., 0., 0., 0., 0., 0., 0.])
self.W0[self.door_index] = self.x_w_cov
def add_models(self, plant, parser, params=None):
iiwa = parser.AddModelFromFile(
"/home/hanikevi/drake/manipulation/models/iiwa_description/sdf/iiwa14_no_collision_no_grav.sdf"
)
plant.WeldFrames(plant.world_frame(),
plant.GetFrameByName("iiwa_link_0"))
#box = Box(10., 10., 10.)
#X_WBox = RigidTransform([0, 0, -5])
#mu = 0.6
#plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box, "ground", CoulombFriction(mu, mu))
#plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box, "ground", [.9, .9, .9, 1.0])
#planar_joint_frame = plant.AddFrame(FixedOffsetFrame("planar_joint_frame", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))
X_WCylinder = RigidTransform([-0.75, 0, 0.5])
hinge = parser.AddModelFromFile(
"/home/hanikevi/drake/examples/manipulation_station/models/simple_hinge.sdf"
)
plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base"),
X_WCylinder)
#cupboard_door_spring = plant.AddForceElement(RevoluteSpring_[float](plant.GetJointByName("right_door_hinge"), nominal_angle = -0.4, stiffness = 10))
if params is None:
bushing = LinearBushingRollPitchYaw_[float](
plant.GetFrameByName("iiwa_link_7"),
plant.GetFrameByName("handle"),
[50, 50, 50], # Torque stiffness
[2., 2., 2.], # Torque damping
[5e4, 5e4, 5e4], # Linear stiffness
[80, 80, 80], # Linear damping
)
else:
print('setting custom stiffnesses')
bushing = LinearBushingRollPitchYaw_[float](
plant.GetFrameByName("iiwa_link_7"),
plant.GetFrameByName("handle"),
[params['k4'], params['k5'], params['k6']], # Torque stiffness
[2, 2, 2], # Torque damping
[params['k1'], params['k2'], params['k3']], # Linear stiffness
[100, 100, 100], # Linear damping
)
bushing_element = plant.AddForceElement(bushing)
return iiwa, hinge, bushing
def cost_stage(self, x, u):
ctrl = 1e-5 * np.sum(u**2)
pos = 15.0 * (x[self.door_index] - self.door_angle_ref)**2
vel = 1e-5 * np.sum(x[8:]**2)
return pos + ctrl + vel
def cost_final(self, x):
return 50 * (1.0 * (x[self.door_index] - self.door_angle_ref)**2 +
np.sum(2.5e-4 * x[8:]**2))
def get_deriv(self, x, u):
self.plant_derivs.SetPositionsAndVelocities(self.plant_derivs_context,
x)
lin = FirstOrderTaylorApproximation(
self.plant_derivs, self.plant_derivs_context,
self.plant.get_actuation_input_port().get_index(),
self.plant.get_state_output_port(self.iiwa).get_index())
return lin.A(), lin.B(), lin.C()
def get_param_deriv(self, x, u):
# Using a closed-form solution; as currently DRAKE doesn't do support autodiff on parameters for LinearBusing.
# Note dC is 0 - stiffness does not affect measurements of q, v
W = self.plant.world_frame()
I = self.plant.GetFrameByName("iiwa_link_7")
H = self.plant.GetFrameByName("handle")
self.plant.SetPositionsAndVelocities(self.plant_context, x)
M = self.plant.CalcMassMatrixViaInverseDynamics(self.plant_context)
Jac_I = self.plant.CalcJacobianSpatialVelocity(
self.plant_context, JacobianWrtVariable.kQDot, I, [0, 0, 0], W, W)
Jac_H = self.plant.CalcJacobianSpatialVelocity(
self.plant_context, JacobianWrtVariable.kQDot, H, [0, 0, 0], W, W)
dA = np.zeros((self.n_x**2, 3))
dA_sq = np.zeros((self.n_x, self.n_x))
for param_ind in range(3):
JH = np.outer(Jac_H[param_ind, :], Jac_H[param_ind, :])
JI = np.outer(Jac_I[param_ind, :], Jac_I[param_ind, :])
dA_sq[8:, :8] = self.dt * np.linalg.inv(M).dot(JH + JI)
dA[:, param_ind] = deepcopy(dA_sq.ravel())
#print(np.sum(np.abs(dA), axis=0))
return dA
def reset(self):
x_traj_new = np.zeros((self.N + 1, self.n_x))
x_traj_new[0, :] = self.x0 + np.multiply(np.sqrt(self.W0),
np.random.randn(self.n_x))
u_traj_new = np.zeros((self.N, self.n_u))
#self.plant_context.SetDiscreteState(x_traj_new[0,:])
self.plant.SetPositionsAndVelocities(self.plant_context,
x_traj_new[0, :])
return x_traj_new, u_traj_new
def rollout(self):
self.u_trj = np.random.randn(self.N, self.n_u) * 0.001
self.x_trj, _ = self.reset()
for i in range(self.N):
self.x_trj[i + 1, :] = self.dyn(self.x_trj[i, :],
self.u_trj[i],
noise=True)
def dyn(self, x, u, phi=None, noise=False):
x_next = self.plant.AllocateDiscreteVariables()
self.plant.SetPositionsAndVelocities(self.plant_context, x)
self.plant.get_actuation_input_port(self.iiwa).FixValue(
self.plant_context, u)
self.plant.CalcDiscreteVariableUpdates(self.plant_context, x_next)
#print(x_next.get_mutable_vector().get_value())
x_new = x_next.get_mutable_vector().get_value()
if noise:
noise = np.multiply(np.sqrt(self.W), np.random.randn(self.n_x))
x_new += noise
return x_new
def obs(self, x, mode=None, noise=False, phi=None):
y = self.plant.get_state_output_port(self.iiwa).Eval(
self.plant_context)
#print(self.bushing.CalcBushingSpatialForceOnFrameA(self.plant_context).translational())
if noise:
y += np.multiply(np.sqrt(self.V), np.random.randn(self.n_y))
return y
def cost(self, x_trj=None, u_trj=None):
cost_trj = 0.0
if x_trj is None:
for i in range(self.N):
cost_trj += self.cost_stage(self.x_trj[i, :], self.u_trj[i, :])
cost_trj += self.cost_final(self.sys.plant, self.x_trj[-1, :])
else:
for i in range(self.N):
cost_trj += self.cost_stage(x_trj[i, :], u_trj[i, :])
cost_trj += self.cost_final(x_trj[-1, :])
return cost_trj
plant_f = MultibodyPlant(0.0)
iiwa_file = FindResourceOrThrow(
"drake/manipulation/models/iiwa_description/sdf/iiwa14_no_collision.sdf")
iiwa = Parser(plant_f).AddModelFromFile(iiwa_file)
# Define some short aliases from frames
W = plant_f.world_frame()
L0 = plant_f.GetFrameByName("iiwa_link_0", iiwa)
L7 = plant_f.GetFrameByName("iiwa_link_7", iiwa)
plant_f.WeldFrames(W, L0)
plant_f.Finalize()
# Allocate float context to be used by evluators.
context_f = plant_f.CreateDefaultContext()
# Create AutoDiffXd plant and corresponding context
plant_ad = plant_f.ToAutoDiffXd()
context_ad = plant_ad.CreateDefaultContext()
def resolve_frame(plant, F):
"""Gets a frame from a plant whose scalar type may be different."""
return plant.GetFrameByName(F.name(), F.model_instance())
# Define target position.
p_WT = [0.1, 0.1, 0.6]
def link_7_distance_to_target(q):
def test_create_q_interpolation(self):
plant = MultibodyPlant(mbp_time_step)
load_atlas(plant, add_ground=False)
context = plant.CreateDefaultContext()
self.skipTest("Unimplemented")
class TestHumanoidPlanner(unittest.TestCase):
def setUp(self):
self.plant = MultibodyPlant(mbp_time_step)
load_atlas(self.plant, add_ground=False)
self.context = self.plant.CreateDefaultContext()
upright_context = self.plant.CreateDefaultContext()
set_atlas_initial_pose(self.plant, upright_context)
self.q_nom = self.plant.GetPositions(upright_context)
self.planner = HumanoidPlanner(self.plant, Atlas.CONTACTS_PER_FRAME, self.q_nom)
self.num_contacts = 16
self.contact_dim = 3*16
self.N_d = 4
def create_default_program(self, N=2):
q_init = self.q_nom.copy()
q_init[6] = 0.94 # Avoid initializing with ground penetration
q_final = q_init.copy()
q_final[4] = 0.0 # x position of pelvis
q_final[6] = 0.93 # z position of pelvis (to make sure final pose touches ground)
num_knot_points = N
max_time = 0.02
# TODO: Use create_minimal_program instead
self.planner.create_full_program(q_init, q_final, num_knot_points, max_time, pelvis_only=True)
def test_getPlantAndContext_float(self):
pelvis_orientation = [1., 0., 0., 0.]
pelvis_position = [5., 6., 7.]
joint_positions = [20 + i for i in range(Atlas.NUM_ACTUATED_DOF)]
q = np.array(pelvis_orientation + pelvis_position + joint_positions)
pelvis_rotational_velocity = [8., 9., 10.]
pelvis_linear_velocity = [11., 12., 13.]
joint_velocity = [100 + i for i in range(Atlas.NUM_ACTUATED_DOF)]
v = np.array(pelvis_rotational_velocity + pelvis_rotational_velocity + joint_velocity)
plant, context = self.planner.getPlantAndContext(q, v)
np.testing.assert_array_almost_equal(plant.GetPositions(context), q)
np.testing.assert_array_almost_equal(plant.GetVelocities(context), v)
def test_getPlantAndContext_AutoDiffXd(self):
pelvis_orientation = [1., 0., 0., 0.]
pelvis_position = [5., 6., 7.]
joint_positions = [20 + i for i in range(Atlas.NUM_ACTUATED_DOF)]
q = np.array(pelvis_orientation + pelvis_position + joint_positions)
q_ad = initializeAutoDiff(q).flatten()
pelvis_rotational_velocity = [8., 9., 10.]
pelvis_linear_velocity = [11., 12., 13.]
joint_velocity = [100 + i for i in range(Atlas.NUM_ACTUATED_DOF)]
v = np.array(pelvis_rotational_velocity + pelvis_rotational_velocity + joint_velocity)
v_ad = initializeAutoDiff(v).flatten()
plant_ad, context_ad = self.planner.getPlantAndContext(q_ad, v_ad)
q_result_ad = plant_ad.GetPositions(context_ad)
v_result_ad = plant_ad.GetVelocities(context_ad)
assert_autodiff_array_almost_equal(q_result_ad, q)
assert_autodiff_array_almost_equal(v_result_ad, v)
def test_reshape_tauj(self):
tau_k = [i for i in range(16)]
tau_j = self.planner.reshape_tauj(tau_k)
tau_j_expected = np.array([[0.0, 0.0, i] for i in range(16)])
np.testing.assert_array_almost_equal(tau_j, tau_j_expected)
def test_get_contact_position(self):
self.skipTest("Unimplemented")
# def test_get_contact_positions_z(self):
# pelvis_orientation = [1., 0., 0., 0.]
# pelvis_position = [0., 0., 0.93845]
# joint_positions = [0.] * Atlas.NUM_ACTUATED_DOF
# q = np.array(pelvis_orientation + pelvis_position + joint_positions)
# pelvis_rotational_velocity = [0., 0., 0.]
# pelvis_linear_velocity = [0., 0., 0.]
# joint_velocity = [0.] * Atlas.NUM_ACTUATED_DOF
# v = np.array(pelvis_rotational_velocity + pelvis_rotational_velocity + joint_velocity)
# contact_positions_z = self.planner.get_contact_positions_z(q, v)
# expected_contact_positions_z = [0.] * Atlas.NUM_CONTACTS
# np.testing.assert_allclose(contact_positions_z, expected_contact_positions_z, atol=epsilon)
# pelvis_orientation = [1., 0., 0., 0.]
# pelvis_position = [0., 0., 1.03845]
# joint_positions = [0.] * Atlas.NUM_ACTUATED_DOF
# q = np.array(pelvis_orientation + pelvis_position + joint_positions)
# pelvis_rotational_velocity = [0., 0., 0.]
# pelvis_linear_velocity = [0., 0., 0.]
# joint_velocity = [0.] * Atlas.NUM_ACTUATED_DOF
# v = np.array(pelvis_rotational_velocity + pelvis_rotational_velocity + joint_velocity)
# contact_positions_z = self.planner.get_contact_positions_z(q, v)
# expected_contact_positions_z = [0.1] * Atlas.NUM_CONTACTS
# np.testing.assert_allclose(contact_positions_z, expected_contact_positions_z, atol=epsilon)
def test_eq7a_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
rdd = np.zeros((N, 3))
rdd[0][2] = -g
rdd[1][2] = -g
self.assertTrue(self.planner.check_eq7a_constraints(F, rdd))
F[0][2::3] = [Atlas.M*g/self.num_contacts]*self.num_contacts
F[1][2::3] = [Atlas.M*g/self.num_contacts]*self.num_contacts
rdd[0][2] = 0.0
rdd[1][2] = 0.0
self.assertTrue(self.planner.check_eq7a_constraints(F, rdd))
def test_eq7b_constraints(self):
self.skipTest("Unimplemented")
def test_eq7c(self):
q = default_q()
v = default_v()
h = np.array([0., 0., 0.])
q_v_h = np.concatenate([q, v, h])
np.testing.assert_allclose(self.planner.eq7c(q_v_h), 0.)
def test_eq7c_constraints(self):
N = 2
self.create_default_program(N)
q = default_q(N)
v = default_v(N)
h = np.array([[0., 0., 0.]]*N)
self.assertTrue(self.planner.check_eq7c_constraints(q, v, h))
def test_eq7d(self):
q = default_q()
qprev = default_q()
w_axis = [0., 0., 1.]
w_mag = [0.0]
v = default_v()
dt = [0.5]
q_qprev_waxis_wmag_v_dt = np.concatenate([q, qprev, w_axis, w_mag, v, dt])
np.testing.assert_allclose(self.planner.eq7d(q_qprev_waxis_wmag_v_dt), 0.)
w_axis = [0.26726124191242438468, 0.53452248382484876937, 0.80178372573727315405]
w_mag = [3.74165738677394138558]
pelvis_rotational_velocity = [1., 2., 3.]
pelvis_linear_velocity = [1., 2., 3.]
joint_velocity = [i for i in range(Atlas.NUM_ACTUATED_DOF)]
v = np.array(pelvis_rotational_velocity + pelvis_linear_velocity + joint_velocity)
pelvis_orientation = [0.5934849924416884, 0.2151038891437094, 0.4302077782874188, 0.6453116674311282]
pelvis_position = [0.5, 1., 1.5+0.93845]
joint_positions = [i*0.5 for i in range(Atlas.NUM_ACTUATED_DOF)]
q = np.array(pelvis_orientation + pelvis_position + joint_positions)
q_qprev_waxis_wmag_v_dt = np.concatenate([q, qprev, w_axis, w_mag, v, dt])
np.testing.assert_allclose(self.planner.eq7d(q_qprev_waxis_wmag_v_dt), 0., atol=epsilon)
def test_eq7d_constraints(self):
N = 2
self.create_default_program(N)
q = default_q(N)
w_axis = np.zeros((N, 3))
w_mag = np.zeros((N,1))
dt = np.zeros((2,1))
w_axis[0] = [0., 0., 1.]
w_mag[0] = 0
pelvis_orientation = [0.5934849924416884, 0.2151038891437094, 0.4302077782874188, 0.6453116674311282]
pelvis_position = [0.5, 1., 1.5+0.93845]
joint_positions = [i*0.5 for i in range(Atlas.NUM_ACTUATED_DOF)]
q[1] = np.array(pelvis_orientation + pelvis_position + joint_positions)
w_axis[1] = [0.26726124191242438468, 0.53452248382484876937, 0.80178372573727315405]
w_mag[1] = 3.74165738677394138558
v = default_v(N)
pelvis_rotational_velocity = [1., 2., 3.]
pelvis_linear_velocity = [1., 2., 3.]
joint_velocity = [i for i in range(Atlas.NUM_ACTUATED_DOF)]
v[1] = np.array(pelvis_rotational_velocity + pelvis_linear_velocity + joint_velocity)
dt[1] = 0.5
self.assertTrue(self.planner.check_eq7d_constraints(q, w_axis, w_mag, v, dt))
def test_eq7e_constraints(self):
N = 2
self.create_default_program(N)
h = np.zeros((N, 3))
hd = np.zeros((N, 3))
dt = np.zeros((N, 1))
dt[1] = 1.0
self.assertTrue(self.planner.check_eq7e_constraints(h, hd, dt))
hd[1] = [1.0, -2.0, 3.0]
h[1] = [0.5, -1.0, 1.5]
dt[1] = 0.5
self.assertTrue(self.planner.check_eq7e_constraints(h, hd, dt))
def test_eq7f_constraints(self):
N = 2
self.create_default_program(N)
r = np.zeros((N, 3))
rd = np.zeros((N, 3))
dt = np.zeros((N, 1))
dt[1] = 1.0
self.assertTrue(self.planner.check_eq7f_constraints(r, rd, dt))
rd[1] = [1.0, -2.0, 3.0]
r[0] = [0.0, 0.0, 0.0]
r[1] = [0.25, -0.5, 0.75]
dt[1] = 0.5
self.assertTrue(self.planner.check_eq7f_constraints(r, rd, dt))
def test_eq7g_constraints(self):
N = 2
self.create_default_program(N)
rd = np.zeros((N, 3))
rdd = np.zeros((N, 3))
dt = np.zeros((N, 1))
dt[1] = 1.0
self.assertTrue(self.planner.check_eq7g_constraints(rd, rdd, dt))
rdd[1] = [1.0, -2.0, 3.0]
rd[1] = [0.5, -1.0, 1.5]
dt[1] = 0.5
self.assertTrue(self.planner.check_eq7g_constraints(rd, rdd, dt))
def test_eq7h(self):
self.skipTest("Unimplemented")
def test_eq7h_constraints(self):
self.skipTest("Unimplemented")
def test_eq7i(self):
self.skipTest("Unimplemented")
def test_eq7i_constraints(self):
self.skipTest("Unimplemented")
def test_eq7j_constraints(self):
N = 2
self.create_default_program(N)
c = np.zeros((N, self.contact_dim))
self.assertTrue(self.planner.check_eq7j_constraints(c))
c[0] = [0.1, -9.9, 5.0] * self.num_contacts
c[1] = [-9.9, 0.1, 0.0] * self.num_contacts
self.assertTrue(self.planner.check_eq7j_constraints(c))
c[0] = [0.1, -9.9, -0.5] * self.num_contacts
c[1] = [-9.9, 0.1, 0.01] * self.num_contacts
self.assertFalse(self.planner.check_eq7j_constraints(c))
def test_eq7k_admissable_posture_constraints(self):
N = 2
self.create_default_program(N)
q = default_q(N)
self.assertTrue(self.planner.check_eq7k_admissable_posture_constraints(q))
q[0, Atlas.FLOATING_BASE_QUAT_DOF + getActuatorIndex(self.plant, "l_leg_kny")] = 0
q[1, Atlas.FLOATING_BASE_QUAT_DOF + getActuatorIndex(self.plant, "l_leg_kny")] = 2.35637
self.assertTrue(self.planner.check_eq7k_admissable_posture_constraints(q))
q[0, Atlas.FLOATING_BASE_QUAT_DOF + getActuatorIndex(self.plant, "l_leg_kny")] = -0.1
q[1, Atlas.FLOATING_BASE_QUAT_DOF + getActuatorIndex(self.plant, "l_leg_kny")] = 2.4
self.assertFalse(self.planner.check_eq7k_admissable_posture_constraints(q))
def test_eq7k_joint_velocity_constraints(self):
N = 2
self.create_default_program(N)
v = default_v(N)
self.assertTrue(self.planner.check_eq7k_joint_velocity_constraints(v))
v[0, Atlas.FLOATING_BASE_DOF + getActuatorIndex(self.plant, "r_leg_kny")] = -12
v[1, Atlas.FLOATING_BASE_DOF + getActuatorIndex(self.plant, "r_leg_kny")] = 12
self.assertTrue(self.planner.check_eq7k_joint_velocity_constraints(v))
v[0, Atlas.FLOATING_BASE_DOF + getActuatorIndex(self.plant, "r_leg_kny")] = -12.1
v[1, Atlas.FLOATING_BASE_DOF + getActuatorIndex(self.plant, "r_leg_kny")] = 12.1
self.assertFalse(self.planner.check_eq7k_joint_velocity_constraints(v))
def test_eq7k_friction_cone_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
beta = np.zeros((N, self.num_contacts*self.N_d))
self.assertTrue(self.planner.check_eq7k_friction_cone_constraints(F, beta))
F[0, 3] = 1.0 # 2nd contact x
F[0, 4] = 2.0 # 2nd contact y
F[0, 5] = 3.0 # 2nd contact z
beta[0, 4] = 1.0 # 2nd contact [1.0, 0.0, 1.0] component
beta[0, 5] = 0.0 # 2nd contact [-1.0, 0.0, 1.0] component
beta[0, 6] = 2.0 # 2nd contact [0.0, 1.0, 1.0] component
beta[0, 7] = 0.0 # 2nd contact [0.0, -1.0, 1.0] component
self.assertTrue(self.planner.check_eq7k_friction_cone_constraints(F, beta))
def test_eq7k_beta_positive_constraints(self):
N = 2
self.create_default_program(N)
beta = np.zeros((N, self.num_contacts*self.N_d))
self.assertTrue(self.planner.check_eq7k_beta_positive_constraints(beta))
beta[1][10] = -0.1
self.assertFalse(self.planner.check_eq7k_beta_positive_constraints(beta))
beta[1][10] = 0.1
self.assertTrue(self.planner.check_eq7k_beta_positive_constraints(beta))
def test_eq7k_torque_constraints(self):
N = 2
self.create_default_program(N)
beta = np.zeros((N, self.num_contacts*self.N_d))
tau = np.zeros((N, self.num_contacts))
self.assertTrue(self.planner.check_eq7k_torque_constraints(tau, beta))
tau[0][3] = 1.0
self.assertFalse(self.planner.check_eq7k_torque_constraints(tau, beta))
tau[0][3] = 0.5
beta[0][4*3] = 0.1
beta[0][4*3+1] = 0.2
beta[0][4*3+2] = 0.3
beta[0][4*3+3] = 0.4
self.assertTrue(self.planner.check_eq7k_torque_constraints(tau, beta))
def test_eq8a_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
c = default_c(N)
slack = default_slack(N)
self.assertTrue(self.planner.check_eq8a_constraints(F, c, slack))
c[1][11*3+2] = 0.0
F[1][11*3+2] = 10.0
self.assertTrue(self.planner.check_eq8a_constraints(F, c, slack))
c[1][10*3+2] = 1.0
F[1][11*3+2] = 10.0
self.assertTrue(self.planner.check_eq8a_constraints(F, c, slack))
c[1][11*3+2] = 1.0
F[1][11*3+2] = 10.0
self.assertFalse(self.planner.check_eq8a_constraints(F, c, slack))
def test_eq8b_constraints(self):
N = 2
self.create_default_program(N)
c = default_c(N)
tau = np.zeros((N, self.num_contacts))
slack = default_slack(N)
self.assertTrue(self.planner.check_eq8b_constraints(tau, c, slack))
c[0][8*3+2] = 0.0
tau[0][8] = -1.0
self.assertTrue(self.planner.check_eq8b_constraints(tau, c, slack))
c[0][9*3+2] = 0.0
tau[0][8] = 1.0
self.assertTrue(self.planner.check_eq8b_constraints(tau, c, slack))
c[0][8*3+2] = 1.0
tau[0][8] = 1.0
self.assertFalse(self.planner.check_eq8b_constraints(tau, c, slack))
def test_eq8c_contact_force_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
self.assertTrue(self.planner.check_eq8c_contact_force_constraints(F))
F[0][3] = -0.1 # 2nd contact x
F[0][4] = -0.1 # 2nd contact y
F[0][5] = 0.1 # 2nd contact z
self.assertTrue(self.planner.check_eq8c_contact_force_constraints(F))
F[0][3] = 0.1 # 2nd contact x
F[0][4] = 0.1 # 2nd contact y
F[0][5] = -0.1 # 2nd contact z
self.assertFalse(self.planner.check_eq8c_contact_force_constraints(F))
def test_eq8c_contact_distance_constraint(self):
N = 2
self.create_default_program(N)
c = default_c(N)
self.assertTrue(self.planner.check_eq8c_contact_distance_constraints(c))
c[1][3*2] = -1.0
c[1][3*2+1] = -1.0
c[1][3*2+2] = 0.1
self.assertTrue(self.planner.check_eq8c_contact_distance_constraints(c))
c[0][3*2] = -1.0
c[0][3*2+1] = -1.0
c[0][3*2+2] = -0.1
self.assertFalse(self.planner.check_eq8c_contact_distance_constraints(c))
def test_eq9a_lhs(self):
self.skipTest("Unimplemented")
def test_eq9a_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
slack = default_slack(N)
self.assertTrue(self.planner.check_eq9a_constraints(F, c, slack))
''' Allow contact point to move if contact force on another contact '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*14] = 0.1 # 15th contact x
c[1, 3*14] = 0.2 # 15th contact x
self.assertTrue(self.planner.check_eq9a_constraints(F, c, slack))
''' Contact point should not move when applying contact force '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*15] = 0.1 # 16th contact x
c[1, 3*15] = 0.2 # 16th contact x
self.assertFalse(self.planner.check_eq9a_constraints(F, c, slack))
''' Allow contact points to move if no contact force '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.0 # 16th contact z
c[0, 3*15] = 0.1 # 16th contact x
c[1, 3*15] = 0.2 # 16th contact x
self.assertTrue(self.planner.check_eq9a_constraints(F, c, slack))
''' This should not check for y axis slipping '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*15+1] = 0.1 # 16th contact y
c[1, 3*15+1] = 0.2 # 16th contact y
self.assertTrue(self.planner.check_eq9a_constraints(F, c, slack))
def test_eq9b_lhs(self):
self.skipTest("Unimplemented")
def test_eq9b_constraints(self):
N = 2
self.create_default_program(N)
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
slack = default_slack(N)
self.assertTrue(self.planner.check_eq9b_constraints(F, c, slack))
''' Allow contact point to move if contact force on another contact '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*14+1] = 0.1 # 15th contact y
c[1, 3*14+1] = 0.2 # 15th contact y
self.assertTrue(self.planner.check_eq9b_constraints(F, c, slack))
''' Contact point should not move when applying contact force '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*15+1] = 0.1 # 16th contact y
c[1, 3*15+1] = 0.2 # 16th contact y
self.assertFalse(self.planner.check_eq9b_constraints(F, c, slack))
''' Allow contact points to move if no contact force '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.0 # 16th contact x
c[0, 3*15+1] = 0.1 # 16th contact y
c[1, 3*15+1] = 0.2 # 16th contact y
self.assertTrue(self.planner.check_eq9b_constraints(F, c, slack))
''' This should not check for x axis slipping '''
F = np.zeros((N, self.contact_dim))
c = np.zeros((N, self.contact_dim))
F[1][3*15 + 2] = 0.2 # 16th contact z
c[0, 3*15] = 0.1 # 16th contact x
c[1, 3*15] = 0.2 # 16th contact x
self.assertTrue(self.planner.check_eq9b_constraints(F, c, slack))
def test_pose_error_cost(self):
self.skipTest("Unimplemented")
def test_standing_trajectory(self):
self.skipTest("Unimplemented")
return
N = 2
t = 0.5
q = np.zeros((N, self.plant.num_positions()))
w_axis = np.zeros((N, 3))
w_axis[:, 2] = 1.0
w_mag = np.zeros((N, 1))
v = np.zeros((N, self.plant.num_velocities()))
dt = t/N*np.ones((N, 1))
r = np.zeros((N, 3))
rd = np.zeros((N, 3))
rdd = np.zeros((N, 3))
c = np.zeros((N, self.contact_dim))
F = np.zeros((N, self.contact_dim))
tau = np.zeros((N, self.num_contacts))
h = np.zeros((N, 3))
hd = np.zeros((N, 3))
beta = np.zeros((N, self.num_contacts*self.N_d))
''' Initialize standing pose '''
for i in range(N):
q[i][0] = 1.0 # w of quaternion
q[i][6] = 0.93845 # z of pelvis
# r[i] = self.planner.calc_r(q[i], v[i])
# TODO
# c[i] = ...
F[i] = np.array([0., 0., Atlas.M*g / 16] * 16)
self.create_default_program()
''' Test all constraints satisfied '''
self.assertTrue(self.planner.check_all_constraints(q, w_axis, w_mag, v, dt, r, rd, rdd, c, F, tau, h, hd, beta))
def test_minimal(self):
self.planner.create_minimal_program(50, 1)
is_success, sol = self.planner.solve()
self.assertTrue(is_success)
def test_0th_order(self):
N = 50
self.planner.create_minimal_program(N, 1)
q_init = default_q()
q_init[6] = 1.5 # z position of pelvis
q_final = default_q()
q_final[0:4] = Quaternion(RollPitchYaw([2*np.pi, np.pi, np.pi/2]).ToRotationMatrix().matrix()).wxyz()
q_final[4] = 1.0 # x position of pelvis
q_final[6] = 2.0 # z position of pelvis
q_final[getJointIndexInGeneralizedPositions(self.planner.plant_float, "r_leg_hpy")] = np.pi/6
self.planner.add_0th_order_constraints(q_init, q_final, False)
is_success, sol = self.planner.solve(self.planner.create_initial_guess())
self.assertTrue(is_success)
visualize(sol.q)
def test_1st_order(self):
N = 50
self.planner.create_minimal_program(N, 0.5)
q_init = default_q()
q_init[6] = 1.5 # z position of pelvis
q_final = default_q()
q_final[0:4] = Quaternion(RollPitchYaw([2*np.pi, np.pi, np.pi/2]).ToRotationMatrix().matrix()).wxyz()
q_final[4] = 1.0 # x position of pelvis
q_final[6] = 2.0 # z position of pelvis
q_final[getJointIndexInGeneralizedPositions(self.planner.plant_float, "r_leg_hpy")] = np.pi/6
self.planner.add_0th_order_constraints(q_init, q_final, False)
self.planner.add_1st_order_constraints()
is_success, sol = self.planner.solve(self.planner.create_initial_guess())
self.assertTrue(is_success)
visualize(sol.q)
pdb.set_trace()
def test_2nd_order(self):
N = 50
self.planner.create_minimal_program(N, 1.0)
q_init = default_q()
q_init[6] = 1.0 # z position of pelvis
q_final = default_q()
q_final[0:4] = Quaternion(RollPitchYaw([0.0, 0.0, np.pi/6]).ToRotationMatrix().matrix()).wxyz()
q_final[4] = 1.0 # x position of pelvis
q_final[6] = 1.5 # z position of pelvis
q_final[getJointIndexInGeneralizedPositions(self.planner.plant_float, "r_leg_hpy")] = np.pi/6
self.planner.add_0th_order_constraints(q_init, q_final, False)
self.planner.add_1st_order_constraints()
self.planner.add_2nd_order_constraints()
is_success, sol = self.planner.solve(self.planner.create_initial_guess())
# if not is_success:
# self.fail("First order solution not found!")
# print("First order solution found!")
# is_success, sol = self.planner.solve(self.planner.create_guess(sol))
self.assertTrue(is_success)
visualize(sol.q)
pdb.set_trace()
def test_contact_sequence_constraints(self):
N = 100
self.planner.create_minimal_program(N, 1.0)
q_init = default_q()
q_init[6] = Atlas.PELVIS_HEIGHT # z position of pelvis
q_final = default_q()
q_final[0:4] = Quaternion(RollPitchYaw([0.0, 0.0, 0.0]).ToRotationMatrix().matrix()).wxyz()
q_final[4] = 0.2 # x position of pelvis
q_final[6] = Atlas.PELVIS_HEIGHT # z position of pelvis
self.planner.add_0th_order_constraints(q_init, q_final, True)
self.planner.add_1st_order_constraints()
self.planner.add_2nd_order_constraints()
is_success, sol = self.planner.solve(self.planner.create_initial_guess(False))
if is_success:
self.planner.add_contact_sequence_constraints()
is_success, sol = self.planner.solve(self.planner.create_guess(sol, False))
visualize(sol.q)
pdb.set_trace()
self.assertTrue(is_success)
def test_complementarity_constraints(self):
N = 50
self.planner.create_minimal_program(N, 1.0)
q_init = default_q()
q_init[6] = Atlas.PELVIS_HEIGHT # z position of pelvis
q_final = default_q()
q_final[0:4] = Quaternion(RollPitchYaw([0.0, 0.0, 0.0]).ToRotationMatrix().matrix()).wxyz()
q_final[4] = 0.5 # x position of pelvis
q_final[6] = Atlas.PELVIS_HEIGHT # z position of pelvis
self.planner.add_0th_order_constraints(q_init, q_final, True)
self.planner.add_1st_order_constraints()
self.planner.add_2nd_order_constraints()
self.planner.add_slack_constraints()
is_success, sol = self.planner.solve(self.planner.create_initial_guess())
if not is_success:
self.fail("Failed to find non-complementarity solution!")
print("Non-complementarity solution found!")
self.planner.add_eq8a_constraints()
self.planner.add_eq8b_constraints()
self.planner.add_eq8c_contact_force_constraints()
self.planner.add_eq8c_contact_distance_constraints()
self.planner.add_eq9a_constraints()
self.planner.add_eq9b_constraints()
is_success, sol = self.planner.solve(self.planner.create_guess(sol))
if not is_success:
self.fail("Failed to find complementarity solution!")
print("Complementarity solution found!")
self.planner.add_slack_cost()
is_success, sol = self.planner.solve(self.planner.create_guess(sol))
# self.planner.add_eq10_cost()
self.assertTrue(is_success)
visualize(sol.q)
pdb.set_trace()
