def _create_patch(self, shape, ax):
type_ = type(shape)
facecolor = shape.facecolor if hasattr(shape, 'facecolor') else (1, 1, 1, 1)
edgecolor = shape.edgecolor if hasattr(shape, 'edgecolor') else (0, 0, 0, 1)
if type_ == pymunk.shapes.Circle:
R = rotation_matrix(shape.body.angle)
p = np.array(shape.body.position) + np.array(shape.offset) @ R.T
r = shape.radius
patch = plt.Circle(p, r, facecolor=facecolor, edgecolor=edgecolor)
ax.add_patch(patch)
elif type_ == pymunk.shapes.Poly:
p = np.array(shape.body.position)
R = rotation_matrix(shape.body.angle)
v = p + np.array(shape.get_vertices()) @ R.T
patch = plt.Polygon(v, facecolor=facecolor, edgecolor=edgecolor, closed=True)
ax.add_patch(patch)
elif type_ == pymunk.shapes.Segment:
p = np.array(shape.body.position)
R = rotation_matrix(shape.body.angle)
v = p + np.array([shape.a, shape.b]) @ R.T
patch = plt.Line2D(v[:, 0], v[:, 1], color=edgecolor, linewidth=1)
ax.add_line(patch)
else:
raise Exception(f'unsupported shape type: {type_}')
self.patches[shape] = patch
def _update_patch(self, shape):
type_ = type(shape)
patch = self.patches[shape]
if type_ == pymunk.shapes.Circle:
R = rotation_matrix(shape.body.angle)
p = np.array(shape.body.position) + np.array(shape.offset) @ R.T
patch.center = p
elif type_ == pymunk.shapes.Poly:
p = np.array(shape.body.position)
R = rotation_matrix(shape.body.angle)
v = p + np.array(shape.get_vertices()) @ R.T
patch.set_xy(v)
elif type_ == pymunk.shapes.Segment:
p = np.array(shape.body.position)
R = rotation_matrix(shape.body.angle)
v = p + np.array([shape.a, shape.b]) @ R.T
patch.set_data(v.T)
else:
raise Exception(f'unsupported shape type: {type_}')
def render(self, ax):
p_ew_w, θ_ew = self.P_ew_w[:2], self.P_ew_w[2]
R_we = util.rotation_matrix(θ_ew)
p_tw_w = p_ew_w + R_we @ self.p_te_e
# sides
p_lw_w = p_tw_w + R_we @ self.p_lt_t
p_rw_w = p_tw_w + R_we @ self.p_rt_t
# contact points
p_c1w_w = p_ew_w + R_we @ self.p_c1e_e
p_c2w_w = p_ew_w + R_we @ self.p_c2e_e
self.tray, = ax.plot([p_lw_w[0], p_rw_w[0]], [p_lw_w[1], p_rw_w[1]],
color='k')
self.com, = ax.plot(p_tw_w[0], p_tw_w[1], 'o', color='k')
self.contacts, = ax.plot([p_c1w_w[0], p_c2w_w[0]],
[p_c1w_w[1], p_c2w_w[1]],
'o',
color='r')
def update_render(self):
p_ew_w, θ_ew = self.P_ew_w[:2], self.P_ew_w[2]
R_we = util.rotation_matrix(θ_ew)
p_tw_w = p_ew_w + R_we @ self.p_te_e
# sides
p_lw_w = p_tw_w + R_we @ self.p_lt_t
p_rw_w = p_tw_w + R_we @ self.p_rt_t
# contact points
p_c1w_w = p_ew_w + R_we @ self.p_c1e_e
p_c2w_w = p_ew_w + R_we @ self.p_c2e_e
self.tray.set_xdata([p_lw_w[0], p_rw_w[0]])
self.tray.set_ydata([p_lw_w[1], p_rw_w[1]])
self.com.set_xdata([p_tw_w[0]])
self.com.set_ydata([p_tw_w[1]])
self.contacts.set_xdata([p_c1w_w[0], p_c2w_w[0]])
self.contacts.set_ydata([p_c1w_w[1], p_c2w_w[1]])
def main():
if TRAY_W < TRAY_MU * TRAY_H:
print('must have w >= μh')
return
N = int(DURATION / SIM_DT) + 1
p_te_e = np.array([0, 0.05 + TRAY_H])
p_oe_e = p_te_e + np.array([0, TRAY_H + OBJ_H])
p_c1e_e = np.array([-TRAY_W, 0.05])
p_c2e_e = np.array([TRAY_W, 0.05])
model = FourInputModel(l1=L1, l2=L2, vel_lim=VEL_LIM, acc_lim=ACC_LIM)
ts = SIM_DT * np.arange(N)
us = np.zeros((N, ni))
P_ew_ws = np.zeros((N, 3))
P_ew_wds = np.zeros((N, 3))
V_ew_ws = np.zeros((N, 3))
P_tw_ws = np.zeros((N, 3))
p_te_es = np.zeros((N, 2))
ineq_cons = np.zeros((N, 3))
p_te_es[0, :] = p_te_e
# state of joints
q0 = np.array([0, 0, 0.25*np.pi, -0.25*np.pi])
dq0 = np.zeros(ni)
X_q = np.concatenate((q0, dq0))
P_ew_w = model.ee_position(X_q)
V_ew_w = model.ee_velocity(X_q)
P_ew_ws[0, :] = P_ew_w
V_ew_ws[0, :] = V_ew_w
# physics simulation
sim = PymunkSimulationTrayBalance(SIM_DT, gravity=-GRAVITY, iterations=30)
sim.add_robot(model, q0, TRAY_W, TRAY_MU)
# tray
tray_body = pymunk.Body(mass=MASS, moment=INERTIA)
tray_body.position = tuple(P_ew_w[:2] + p_te_e)
tray_corners = [(-RADIUS, TRAY_H), (-RADIUS, -TRAY_H), (RADIUS, -TRAY_H),
(RADIUS, TRAY_H)]
tray = pymunk.Poly(tray_body, tray_corners, radius=0)
tray.facecolor = (0.25, 0.5, 1, 1)
tray.friction = TRAY_MU
tray.collision_type = 1
sim.space.add(tray.body, tray)
obj_body = pymunk.Body(mass=OBJ_MASS, moment=OBJ_INERTIA)
obj_body.position = tuple(P_ew_w[:2] + p_oe_e)
obj_corners = [(-OBJ_W, OBJ_H), (-OBJ_W, -OBJ_H), (OBJ_W, -OBJ_H),
(OBJ_W, OBJ_H)]
obj = pymunk.Poly(obj_body, obj_corners, radius=0)
obj.facecolor = (0.5, 0.5, 0.5, 1)
obj.friction = OBJ_MU / TRAY_MU # so that mu with tray = OBJ_MU
obj.collision_type = 1
# sim.space.add(obj.body, obj)
# sim.space.add_collision_handler(1, 1).post_solve = tray_cb
# reference trajectory
# timescaling = trajectories.QuinticTimeScaling(DURATION)
# trajectory = trajectories.Circle(P_ew_w[:2], 0.25, timescaling, DURATION)
trajectory = trajectories.Point(P_ew_w[:2] + np.array([1, -1]))
P_ew_wds[:, :2], _, _ = trajectory.sample(ts)
# P_ew_wds[:, 2] = -np.pi / 2
# rendering
goal_renderer = plotting.PointRenderer(P_ew_wds[-1, :2], color='r')
sim_renderer = PymunkRenderer(sim.space, sim.markers)
# trajectory_renderer = plotting.TrajectoryRenderer(trajectory, ts)
renderers = [goal_renderer, sim_renderer]
video = plotting.Video(name='tray_balance.mp4', fps=1./(SIM_DT*PLOT_PERIOD))
plotter = plotting.RealtimePlotter(renderers, video=video)
plotter.start() # TODO for some reason setting grid=True messes up the base rendering
# pymunk rendering
# plt.ion()
# fig = plt.figure()
# ax = plt.gca()
# plt.grid()
#
# ax.set_xlabel('x (m)')
# ax.set_ylabel('y (m)')
# ax.set_xlim([-1, 6])
# ax.set_ylim([-1, 2])
#
# ax.set_aspect('equal')
#
# options = pymunk.matplotlib_util.DrawOptions(ax)
# options.flags = pymunk.SpaceDebugDrawOptions.DRAW_SHAPES
def objective_unrolled(X_q_0, X_ee_d, var):
''' Unroll the objective over n timesteps. '''
obj = 0
X_q = X_q_0
for i in range(MPC_STEPS):
u = var[i*nv:(i+1)*nv]
X_q = model.step_unconstrained(X_q, u, MPC_DT)
X_ee = model.ee_state(X_q)
e = X_ee_d[i*ns_ee:(i+1)*ns_ee] - X_ee
obj = obj + 0.5 * (e @ W @ e + X_q @ Q @ X_q + u @ R @ u)
return obj
def ineq_constraints(X_ee, a_ee, jnp=jnp):
θ_ew, dθ_ew = X_ee[2], X_ee[5]
a_ew_w, ddθ_ew = a_ee[:2], a_ee[2]
R_ew = jnp.array([[ jnp.cos(θ_ew), jnp.sin(θ_ew)],
[-jnp.sin(θ_ew), jnp.cos(θ_ew)]])
S1 = skew1(1)
g = jnp.array([0, GRAVITY])
ddR_we = (ddθ_ew*S1 - dθ_ew**2*jnp.eye(2)) @ R_ew.T
α1, α2 = MASS * R_ew @ (a_ew_w+g) + MASS * (ddθ_ew*S1 - dθ_ew**2*jnp.eye(2)) @ p_te_e
# α1, α2 = OBJ_MASS * R_ew @ (a_ew_w + g) + OBJ_MASS * (ddθ_ew*S1 - dθ_ew**2*jnp.eye(2)) @ p_oe_e
α3 = INERTIA * ddθ_ew
# constraints considering y acceleration at each contact point
# β1 = jnp.array([0, 1]) @ R_ew @ (a_ew_w + ddR_we @ p_c1e_e + g)
# β2 = jnp.array([0, 1]) @ R_ew @ (a_ew_w + ddR_we @ p_c2e_e + g)
# ineq_con = jnp.array([
# TRAY_MU*jnp.abs(α2) - jnp.abs(α1),
# β1,
# β2])
h1 = TRAY_MU*jnp.abs(α2) - jnp.abs(α1)
# h1 = OBJ_MU*jnp.abs(α2) - jnp.abs(α1)
h2 = α2
# h3 = TRAY_H**2*α1**2 + TRAY_W*(TRAY_W-2*TRAY_MU*TRAY_H)*α2**2 - α3**2
# h3 = OBJ_H**2 * α1**2 + OBJ_W * (OBJ_W - 2*OBJ_MU*OBJ_H)*α2**2 - α3**2
h3 = 1
return jnp.array([h1, h2, h3])
def ineq_constraints_unrolled(X_q_0, X_ee_d, var):
# var is now just the lifted joint acceleration inputs
X_q = X_q_0
ineq_con = jnp.zeros(MPC_STEPS * nc_ineq)
for i in range(MPC_STEPS):
u = var[i*nv:(i+1)*nv]
X_ee = model.ee_state(X_q)
a_ee = model.ee_acceleration(X_q, u)
ineq_coni = ineq_constraints(X_ee, a_ee)
ineq_con = jax.ops.index_update(ineq_con, jax.ops.index[i*nc_ineq:(i+1)*nc_ineq], ineq_coni)
X_q = model.step_unconstrained(X_q, u, MPC_DT)
return ineq_con
# X_ee = model.ee_state(X_q)
# a_ee = model.ee_acceleration(X_q, np.zeros(4))
# ineq_constraints(X_ee, a_ee)
# return
# Construct the SQP controller
obj_fun = jax.jit(objective_unrolled)
obj_jac = jax.jit(jax.jacfwd(objective_unrolled, argnums=2))
obj_hess = jax.jit(hessian(objective_unrolled, argnums=2))
objective = sqp.Objective(obj_fun, obj_jac, obj_hess)
lbA = np.zeros(MPC_STEPS * nc)
ubA = np.infty * np.ones(MPC_STEPS * nc)
con_fun = jax.jit(ineq_constraints_unrolled)
con_jac = jax.jit(jax.jacfwd(ineq_constraints_unrolled, argnums=2))
constraints = sqp.Constraints(con_fun, con_jac, lbA, ubA)
lb = -ACC_LIM * np.ones(MPC_STEPS * nv)
ub = ACC_LIM * np.ones(MPC_STEPS * nv)
bounds = sqp.Bounds(lb, ub)
controller = sqp.SQP(nv*MPC_STEPS, nc*MPC_STEPS, objective, constraints, bounds, num_iter=3, verbose=True)
for i in range(N - 1):
t = ts[i+1]
if i % CTRL_PERIOD == 0:
t_sample = np.minimum(t + MPC_DT*np.arange(MPC_STEPS), DURATION)
pd, vd, _ = trajectory.sample(t_sample, flatten=False)
z = np.zeros((MPC_STEPS, 1))
X_ee_d = np.hstack((pd, z, vd, z)).flatten()
# -0.5*np.pi*np.ones((MPC_STEPS, 1))
# start = time.time()
var = controller.solve(X_q, X_ee_d)
# print(time.time() - start)
u = var[:ni] # joint acceleration input
sim.command_acceleration(u)
# integrate the system
X_q = np.concatenate(sim.step())
P_ew_w = model.ee_position(X_q)
V_ew_w = model.ee_velocity(X_q)
X_ee = model.ee_state(X_q)
a_ee = model.ee_acceleration(X_q, u)
ineq_cons[i+1, :] = ineq_constraints(X_ee, a_ee, jnp=np)
# if (ineq_con < 0).any():
# IPython.embed()
# tray position is (ideally) a constant offset from EE frame
θ_ew = P_ew_w[2]
R_we = util.rotation_matrix(θ_ew)
P_tw_w = np.array([tray.body.position.x, tray.body.position.y, tray.body.angle])
# record
us[i, :] = u
P_ew_ws[i+1, :] = P_ew_w
V_ew_ws[i+1, :] = V_ew_w
P_tw_ws[i+1, :] = P_tw_w
p_te_es[i+1, :] = R_we.T @ (P_tw_w[:2] - P_ew_w[:2])
if i % PLOT_PERIOD == 0:
# ax.cla()
# ax.set_xlim([-1, 6])
# ax.set_ylim([-1, 2])
#
# sim.space.debug_draw(options)
# fig.canvas.draw()
# fig.canvas.flush_events()
plotter.update()
plotter.done()
v_te_es = (p_te_es[1:] - p_te_es[:-1]) / SIM_DT
v_te_es_smooth = np.zeros_like(v_te_es)
v_te_es_smooth[:, 0] = np.convolve(v_te_es[:, 0], np.ones(100) / 100, 'same')
v_te_es_smooth[:, 1] = np.convolve(v_te_es[:, 1], np.ones(100) / 100, 'same')
plt.figure()
plt.plot(ts[:N], P_ew_wds[:, 0], label='$x_d$', color='b', linestyle='--')
plt.plot(ts[:N], P_ew_wds[:, 1], label='$y_d$', color='r', linestyle='--')
plt.plot(ts[:N], P_ew_ws[:, 0], label='$x$', color='b')
plt.plot(ts[:N], P_ew_ws[:, 1], label='$y$', color='r')
plt.plot(ts[:N], P_tw_ws[:, 0], label='$t_x$')
plt.plot(ts[:N], P_tw_ws[:, 1], label='$t_y$')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('Position')
plt.title('End effector position')
plt.figure()
plt.plot(ts[:N], p_te_es[:, 0], label='$x$', color='b')
plt.plot(ts[:N], p_te_es[:, 1], label='$y$', color='r')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('$p^{te}_e$')
plt.title('$p^{te}_e$')
plt.figure()
plt.plot(ts[:N], p_te_e[0] - p_te_es[:, 0], label='$x$', color='b')
plt.plot(ts[:N], p_te_e[1] - p_te_es[:, 1], label='$y$', color='r')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('$p^{te}_e$ error')
plt.title('$p^{te}_e$ error')
# plt.figure()
# plt.plot(ts[:N-1], v_te_es[:, 0], label='$v_x$', color='b')
# plt.plot(ts[:N-1], v_te_es[:, 1], label='$v_y$', color='r')
# plt.plot(ts[:N-1], v_te_es_smooth[:, 0], label='$v_x$ (smooth)')
# plt.plot(ts[:N-1], v_te_es_smooth[:, 1], label='$v_y$ (smooth)')
# plt.grid()
# plt.legend()
# plt.xlabel('Time (s)')
# plt.ylabel('$v^{te}_e$')
# plt.title('$v^{te}_e$')
plt.figure()
plt.plot(ts[:N], ineq_cons[:, 0], label='$h_1$')
plt.plot(ts[:N], ineq_cons[:, 1], label='$h_2$')
plt.plot(ts[:N], ineq_cons[:, 2], label='$h_3$')
plt.plot([0, ts[-1]], [0, 0], color='k')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('Value')
plt.title('Inequality constraints')
plt.show()
def update_render(self):
v = self.vertices @ rotation_matrix(self.angle).T
self.patch.set_xy(self.p + v)
def render(self, ax):
v = self.vertices @ rotation_matrix(self.angle).T
self.patch = plt.Polygon(self.p + v, color='k', closed=True)
ax.add_patch(self.patch)
def get_position(self):
R = rotation_matrix(self.body.angle)
return np.array(self.body.position) + self.offset @ R.T
def main():
if TRAY_W < TRAY_MU * TRAY_H:
print('must have w >= μh')
return
N = int(DURATION / SIM_DT) + 1
p_te_e = np.array([0, 0.05 + TRAY_H])
p_oe_e = p_te_e + np.array([0, TRAY_H + OBJ_H])
r_c1e_e = np.array([-TRAY_W, 0.05])
r_c2e_e = np.array([TRAY_W, 0.05])
r_tc1_t = np.array([TRAY_W, TRAY_H])
r_tc2_t = np.array([-TRAY_W, TRAY_H])
r_t2t1_t = -r_tc2_t + r_tc1_t
model = FourInputModel(l1=L1, l2=L2, vel_lim=VEL_LIM, acc_lim=ACC_LIM)
ts = SIM_DT * np.arange(N)
us = np.zeros((N, ni))
P_ew_ws = np.zeros((N, 3))
P_ew_wds = np.zeros((N, 3))
V_ew_ws = np.zeros((N, 3))
P_tw_ws = np.zeros((N, 3))
p_te_es = np.zeros((N, 2))
ineq_cons = np.zeros((N, 3))
p_te_es[0, :] = p_te_e
# state of joints
q0 = np.array([0, 0, 0.25 * np.pi, -0.25 * np.pi])
dq0 = np.zeros(ni)
X_q = np.concatenate((q0, dq0))
P_ew_w = model.ee_position(X_q)
V_ew_w = model.ee_velocity(X_q)
P_ew_ws[0, :] = P_ew_w
V_ew_ws[0, :] = V_ew_w
# physics simulation
sim = PymunkSimulationTrayBalance(SIM_DT, gravity=-GRAVITY, iterations=30)
sim.add_robot(model, q0, TRAY_W, TRAY_MU)
# tray
tray_body = pymunk.Body(mass=MASS, moment=INERTIA)
tray_body.position = tuple(P_ew_w[:2] + p_te_e)
tray_corners = [(-RADIUS, TRAY_H), (-RADIUS, -TRAY_H), (RADIUS, -TRAY_H),
(RADIUS, TRAY_H)]
tray = pymunk.Poly(tray_body, tray_corners, radius=0)
tray.facecolor = (0.25, 0.5, 1, 1)
tray.friction = TRAY_MU
tray.collision_type = 1
sim.space.add(tray.body, tray)
obj_body = pymunk.Body(mass=OBJ_MASS, moment=OBJ_INERTIA)
obj_body.position = tuple(P_ew_w[:2] + p_oe_e)
obj_corners = [(-OBJ_W, OBJ_H), (-OBJ_W, -OBJ_H), (OBJ_W, -OBJ_H),
(OBJ_W, OBJ_H)]
obj = pymunk.Poly(obj_body, obj_corners, radius=0)
obj.facecolor = (0.5, 0.5, 0.5, 1)
obj.friction = OBJ_MU / TRAY_MU # so that mu with tray = OBJ_MU
obj.collision_type = 1
# sim.space.add(obj.body, obj)
# sim.space.add_collision_handler(1, 1).post_solve = tray_cb
# reference trajectory
# timescaling = trajectories.QuinticTimeScaling(DURATION)
# trajectory = trajectories.Circle(P_ew_w[:2], 0.25, timescaling, DURATION)
trajectory = trajectories.Point(P_ew_w[:2] + np.array([1, -1]))
P_ew_wds[:, :2], _, _ = trajectory.sample(ts)
# P_ew_wds[:, 2] = -np.pi / 2
# rendering
goal_renderer = plotting.PointRenderer(P_ew_wds[-1, :2], color='r')
sim_renderer = PymunkRenderer(sim.space, sim.markers)
# trajectory_renderer = plotting.TrajectoryRenderer(trajectory, ts)
renderers = [goal_renderer, sim_renderer]
video = plotting.Video(name='tray_balance.mp4',
fps=1. / (SIM_DT * PLOT_PERIOD))
plotter = plotting.RealtimePlotter(renderers, video=video)
plotter.start(
) # TODO for some reason setting grid=True messes up the base rendering
# pymunk rendering
# plt.ion()
# fig = plt.figure()
# ax = plt.gca()
# plt.grid()
#
# ax.set_xlabel('x (m)')
# ax.set_ylabel('y (m)')
# ax.set_xlim([-1, 6])
# ax.set_ylim([-1, 2])
#
# ax.set_aspect('equal')
#
# options = pymunk.matplotlib_util.DrawOptions(ax)
# options.flags = pymunk.SpaceDebugDrawOptions.DRAW_SHAPES
def jax_rotation_matrix(θ):
return jnp.array([[jnp.cos(θ), -jnp.sin(θ)], [jnp.sin(θ), jnp.cos(θ)]])
def calc_tray_pose(q):
"""Calculate tray pose in terms of x1, y1, y2."""
x1, y1, y2 = q
Δy = y2 - y1
Δx = jnp.sqrt((2 * TRAY_W)**2 - Δy**2)
θ_tw = jnp.arctan2(Δy, Δx)
R_wt = jax_rotation_matrix(θ_tw)
r_c1w_w = jnp.array([x1, y1])
r_tw_w = r_c1w_w + R_wt @ r_tc1_t
return jnp.array([r_tw_w[0], r_tw_w[1], θ_tw])
# def calc_tray_state(P_tw_w, V_tw_w):
# """Calculate tray parameters q given the pose."""
# # TODO ideally we'd just keep everything in terms of q, but the old
# # approach parameterizes the tray by its pose.
# # TODO inconsistent notation
# r_tw_w, θ_tw = P_tw_w[:2], P_tw_w[2]
# R_wt = jnp.array([[jnp.cos(θ_tw), -jnp.sin(θ_tw)],
# [jnp.sin(θ_tw), jnp.cos(θ_tw)]])
# r_t1w_w = r_tw_w - R_wt @ r_tc1_t
# r_t2w_w = r_tw_w - R_wt @ r_tc2_t
# q = jnp.array([r_t1w_w[0], r_t1w_w[1], r_t2w_w[1]])
#
# return q, dq
calc_tray_jacobian = jax.jit(jax.jacfwd(calc_tray_pose))
def calc_tray_potential(q):
"""Calculate tray potential energy."""
x, y, θ = calc_tray_pose(q)
return MASS * GRAVITY * y
calc_g = jax.jit(jax.jacfwd(calc_tray_potential))
def calc_tray_mass_matrix(q):
"""Calculate tray mass matrix."""
J = calc_tray_jacobian(q)
M = J.T @ TRAY_G_MAT @ J
return M
def calc_ddR_wt(q, dq, ddq):
"""Calculate the second time derivative of tray's rotation matrix."""
Δy = q[2] - q[1]
dΔy = dq[2] - dq[1]
ddΔy = ddq[2] - ddq[1]
Δx = jnp.sqrt((2 * TRAY_W)**2 - Δy**2)
θ_tw = jnp.arctan2(Δy, Δx)
dθ_tw = dΔy / Δx # TODO possible division by zero
ddθ_tw = (ddΔy * Δx**2 + dΔy**2 * Δy) / Δx**3
S1 = skew1(1)
ddR_wt = (ddθ_tw * S1 - dθ_tw**2) @ jax_rotation_matrix(θ_tw)
return ddR_wt
calc_dMdq = jax.jit(jax.jacfwd(calc_tray_mass_matrix))
def calc_tray_h(q, dq):
"""Calculate h term in dynamics of the form M(q)*ddq + h(q,dq) = f."""
dMdq = calc_dMdq(q)
Γ = dMdq - 0.5 * dMdq.T
g = calc_g(q)
h = dq.dot(Γ).dot(dq) + g
return h
def objective_unrolled(X_q_0, X_ee_d, var):
"""Unroll the objective over n timesteps."""
obj = 0
X_q = X_q_0
for i in range(MPC_STEPS):
u = var[i * nv:(i + 1) * nv]
X_q = model.step_unconstrained(X_q, u, MPC_DT)
X_ee = model.ee_state(X_q)
e = X_ee_d[i * ns_ee:(i + 1) * ns_ee] - X_ee
obj = obj + 0.5 * (e @ W @ e + X_q @ Q @ X_q + u @ R @ u)
return obj
def ineq_constraints(X_ee, a_ee, jnp=jnp):
r_ew_w = X_ee[:2]
v_ew_w = X_ee[3:5]
θ_ew, dθ_ew = X_ee[2], X_ee[5]
a_ew_w, ddθ_ew = a_ee[:2], a_ee[2]
R_ew = jnp.array([[jnp.cos(θ_ew), jnp.sin(θ_ew)],
[-jnp.sin(θ_ew), jnp.cos(θ_ew)]])
S1 = skew1(1)
g = jnp.array([0, GRAVITY])
# calculate tray state from EE state
r_c1w_w = r_ew_w + R_ew.T @ r_c1e_e
r_c2w_w = r_ew_w + R_ew.T @ r_c2e_e
v_c1w_w = v_ew_w + dθ_ew * S1 @ R_ew.T @ r_c1e_e
v_c2w_w = v_ew_w + dθ_ew * S1 @ R_ew.T @ r_c2e_e
qt = jnp.array([r_c1w_w[0], r_c1w_w[1], r_c2w_w[1]])
dqt = jnp.array([v_c1w_w[0], v_c1w_w[1], v_c2w_w[1]])
# calculate tray dynamics
Mt = calc_tray_mass_matrix(qt)
ht = calc_tray_h(qt, dqt)
Qt = jnp.linalg.solve(
Mt, ht) # TODO bad name: this isn't the generalized forces
ddR_we = (ddθ_ew * S1 - dθ_ew**2 * jnp.eye(2)) @ R_ew.T
ddR_wt = calc_ddR_wt(qt, dqt, Qt)
a_t1w_w = Qt[:2]
a_t2w_w = a_t1w_w + ddR_wt @ r_t2t1_t
ddy1 = np.array([0, 1]) @ R_ew @ (a_ew_w + ddR_we @ r_c1e_e)
ddy2 = np.array([0, 1]) @ R_ew @ (a_ew_w + ddR_we @ r_c2e_e)
# TODO: how to deal with friction under this formulation?
h1 = 1
h2 = ddy1 + a_t1w_w[1]
h3 = ddy2 + a_t2w_w[1]
# α1, α2 = MASS * R_ew @ (a_ew_w+g) + MASS * (ddθ_ew*S1 - dθ_ew**2*jnp.eye(2)) @ p_te_e
# α3 = INERTIA * ddθ_ew
#
# h1 = TRAY_MU*jnp.abs(α2) - jnp.abs(α1)
# h2 = α2
# h3 = 1
return jnp.array([h1, h2, h3])
def ineq_constraints_unrolled(X_q_0, X_ee_d, var):
# var is now just the lifted joint acceleration inputs
X_q = X_q_0
ineq_con = jnp.zeros(MPC_STEPS * nc_ineq)
for i in range(MPC_STEPS):
u = var[i * nv:(i + 1) * nv]
X_ee = model.ee_state(X_q)
a_ee = model.ee_acceleration(X_q, u)
ineq_coni = ineq_constraints(X_ee, a_ee)
ineq_con = jax.ops.index_update(
ineq_con, jax.ops.index[i * nc_ineq:(i + 1) * nc_ineq],
ineq_coni)
X_q = model.step_unconstrained(X_q, u, MPC_DT)
return ineq_con
# Construct the SQP controller
obj_fun = jax.jit(objective_unrolled)
obj_jac = jax.jit(jax.jacfwd(objective_unrolled, argnums=2))
obj_hess = jax.jit(hessian(objective_unrolled, argnums=2))
objective = sqp.Objective(obj_fun, obj_jac, obj_hess)
lbA = np.zeros(MPC_STEPS * nc)
ubA = np.infty * np.ones(MPC_STEPS * nc)
con_fun = jax.jit(ineq_constraints_unrolled)
con_jac = jax.jit(jax.jacfwd(ineq_constraints_unrolled, argnums=2))
constraints = sqp.Constraints(con_fun, con_jac, lbA, ubA)
lb = -ACC_LIM * np.ones(MPC_STEPS * nv)
ub = ACC_LIM * np.ones(MPC_STEPS * nv)
bounds = sqp.Bounds(lb, ub)
controller = sqp.SQP(nv * MPC_STEPS,
nc * MPC_STEPS,
objective,
constraints,
bounds,
num_iter=3,
verbose=False)
for i in range(N - 1):
t = ts[i + 1]
if i % CTRL_PERIOD == 0:
t_sample = np.minimum(t + MPC_DT * np.arange(MPC_STEPS), DURATION)
pd, vd, _ = trajectory.sample(t_sample, flatten=False)
z = np.zeros((MPC_STEPS, 1))
X_ee_d = np.hstack((pd, z, vd, z)).flatten()
# -0.5*np.pi*np.ones((MPC_STEPS, 1))
# start = time.time()
var = controller.solve(X_q, X_ee_d)
# print(time.time() - start)
u = var[:ni] # joint acceleration input
sim.command_acceleration(u)
# integrate the system
X_q = np.concatenate(sim.step())
P_ew_w = model.ee_position(X_q)
V_ew_w = model.ee_velocity(X_q)
X_ee = model.ee_state(X_q)
a_ee = model.ee_acceleration(X_q, u)
ineq_cons[i + 1, :] = ineq_constraints(X_ee, a_ee, jnp=np)
# if (ineq_con < 0).any():
# IPython.embed()
# tray position is (ideally) a constant offset from EE frame
θ_ew = P_ew_w[2]
R_we = util.rotation_matrix(θ_ew)
P_tw_w = np.array(
[tray.body.position.x, tray.body.position.y, tray.body.angle])
# record
us[i, :] = u
P_ew_ws[i + 1, :] = P_ew_w
V_ew_ws[i + 1, :] = V_ew_w
P_tw_ws[i + 1, :] = P_tw_w
p_te_es[i + 1, :] = R_we.T @ (P_tw_w[:2] - P_ew_w[:2])
if i % PLOT_PERIOD == 0:
# ax.cla()
# ax.set_xlim([-1, 6])
# ax.set_ylim([-1, 2])
#
# sim.space.debug_draw(options)
# fig.canvas.draw()
# fig.canvas.flush_events()
plotter.update()
plotter.done()
v_te_es = (p_te_es[1:] - p_te_es[:-1]) / SIM_DT
v_te_es_smooth = np.zeros_like(v_te_es)
v_te_es_smooth[:, 0] = np.convolve(v_te_es[:, 0],
np.ones(100) / 100, 'same')
v_te_es_smooth[:, 1] = np.convolve(v_te_es[:, 1],
np.ones(100) / 100, 'same')
plt.figure()
plt.plot(ts[:N], P_ew_wds[:, 0], label='$x_d$', color='b', linestyle='--')
plt.plot(ts[:N], P_ew_wds[:, 1], label='$y_d$', color='r', linestyle='--')
plt.plot(ts[:N], P_ew_ws[:, 0], label='$x$', color='b')
plt.plot(ts[:N], P_ew_ws[:, 1], label='$y$', color='r')
plt.plot(ts[:N], P_tw_ws[:, 0], label='$t_x$')
plt.plot(ts[:N], P_tw_ws[:, 1], label='$t_y$')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('Position')
plt.title('End effector position')
plt.figure()
plt.plot(ts[:N], p_te_es[:, 0], label='$x$', color='b')
plt.plot(ts[:N], p_te_es[:, 1], label='$y$', color='r')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('$p^{te}_e$')
plt.title('$p^{te}_e$')
plt.figure()
plt.plot(ts[:N], p_te_e[0] - p_te_es[:, 0], label='$x$', color='b')
plt.plot(ts[:N], p_te_e[1] - p_te_es[:, 1], label='$y$', color='r')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('$p^{te}_e$ error')
plt.title('$p^{te}_e$ error')
# plt.figure()
# plt.plot(ts[:N-1], v_te_es[:, 0], label='$v_x$', color='b')
# plt.plot(ts[:N-1], v_te_es[:, 1], label='$v_y$', color='r')
# plt.plot(ts[:N-1], v_te_es_smooth[:, 0], label='$v_x$ (smooth)')
# plt.plot(ts[:N-1], v_te_es_smooth[:, 1], label='$v_y$ (smooth)')
# plt.grid()
# plt.legend()
# plt.xlabel('Time (s)')
# plt.ylabel('$v^{te}_e$')
# plt.title('$v^{te}_e$')
plt.figure()
plt.plot(ts[:N], ineq_cons[:, 0], label='$h_1$')
plt.plot(ts[:N], ineq_cons[:, 1], label='$h_2$')
plt.plot(ts[:N], ineq_cons[:, 2], label='$h_3$')
plt.plot([0, ts[-1]], [0, 0], color='k')
plt.grid()
plt.legend()
plt.xlabel('Time (s)')
plt.ylabel('Value')
plt.title('Inequality constraints')
plt.show()
def _lookahead(self, q0, pr, u, N, pc):
''' Generate lifted matrices proprogating the state N timesteps into the
future. '''
ni = self.model.ni # number of joints
no = self.model.no # number of Cartesian outputs
fbar = np.zeros(no * N) # Lifted forward kinematics
Jbar = np.zeros((no * N, ni * N)) # Lifted Jacobian
Qbar = np.kron(np.eye(N), self.Q)
Rbar = np.kron(np.eye(N), self.R)
# lower triangular matrix of ni*ni identity matrices
Ebar = np.kron(np.tril(np.ones((N, N))), np.eye(ni))
# Integrate joint positions from the last iteration
qbar = np.tile(q0, N + 1)
qbar[ni:] = qbar[ni:] + self.dt * Ebar.dot(u)
# TODO: need to integrate pc as well: this takes the place of fbar
num_body_pts = 2 + 1
Abar = np.zeros((N * num_body_pts, ni * N))
lbA = np.zeros(N * num_body_pts)
for k in range(N):
q = qbar[(k + 1) * ni:(k + 2) * ni]
# calculate points defining front of base
pb = q[:2]
θb = q[2]
R = util.rotation_matrix(θb)
rx = 0.5
ry = 0.25
p1 = R.dot(np.array([rx, ry]))
p2 = R.dot(np.array([rx, -ry]))
# pf is the closest point to the line segment
pf, _ = util.dist_to_line_segment(pc, p1, p2)
# transform into body frame
b_pf = R.T.dot(pf - pb)
JR = util.rotation_jacobian(θb)
Jf = np.hstack((R, JR.dot(pb + b_pf)[:, None], np.zeros((2, 2))))
pe = self.model.forward(q)
Je = self.model.jacobian(q)
re = (pc - pe) / np.linalg.norm(pc - pe)
rf = (pc - pf) / np.linalg.norm(pc - pf)
# propagate center of object forward
pc = pc + self.dt * (Jf + Je).dot(u)
# fbar[k*no:(k+1)*no] = pm
# Jbar[k*no:(k+1)*no, k*ni:(k+1)*ni] = Jm
# # TODO hardcoded radius
# # EE and obstacle
# d_ee_obs = np.linalg.norm(p - pc) - 0.45
# Abar[k*num_body_pts, k*ni:(k+1)*ni] = (p - pc).T.dot(J) / np.linalg.norm(p - pc)
# lbA[k*num_body_pts] = -d_ee_obs
#
# # base and obstacle
# pb = q[:2]
# Jb = np.array([[1, 0, 0, 0, 0],
# [0, 1, 0, 0, 0]])
# d_base_obs = np.linalg.norm(pb - pc) - 0.45 - 0.56
# Abar[k*num_body_pts+1, k*ni:(k+1)*ni] = (pb - pc).T.dot(Jb) / np.linalg.norm(pb - pc)
# lbA[k*num_body_pts+1] = -d_base_obs
#
# # pf and ee: these need to stay close together
# pf = self.model.forward_f(q)
# Jf = self.model.jacobian_f(q)
# d_pf_ee = np.linalg.norm(p - pf)
# A_pf_ee = -(pf - p).T.dot(Jf - J) / d_pf_ee
# lbA_pf_ee = d_pf_ee - 0.75
# Abar[k*num_body_pts+2, k*ni:(k+1)*ni] = A_pf_ee
# lbA[k*num_body_pts+2] = lbA_pf_ee
dbar = fbar - pr
H = Rbar + self.dt**2 * Ebar.T.dot(
Jbar.T).dot(Qbar).dot(Jbar).dot(Ebar)
g = u.T.dot(Rbar) + self.dt * dbar.T.dot(Qbar).dot(Jbar).dot(Ebar)
A = self.dt * Abar.dot(Ebar)
return H, g, A, lbA