class CollectDocking:
def __init__(self, figure, ax, num_data, bfname, traj_dir):
self.figure = figure
self.ax = ax
self.press = False
self.vmp = VMP(dim=2, kernel_num=10)
self.bfname = bfname
self.traj_dir = traj_dir
if not os.path.exists(traj_dir):
try:
os.makedirs(traj_dir)
except OSError:
raise
self.id = 0
self.num_data = num_data
self.trqueries = np.random.uniform(low=(0, 0, 0, 10, 10, 0),
high=(10, 10, 2 * np.pi, 20, 20,
2 * np.pi),
size=(num_data, 6))
obsExp = ObsExp(exp_name="Docking")
self.envs = obsExp.get_envs(self.trqueries)
self.next()
self.connect()
def connect(self):
'connect to all the events we need'
self.cidpress = self.figure.canvas.mpl_connect('button_press_event',
self.on_press)
self.cidrelease = self.figure.canvas.mpl_connect(
'button_release_event', self.on_release)
self.cidmotion = self.figure.canvas.mpl_connect(
'motion_notify_event', self.on_motion)
self.cidkey = self.figure.canvas.mpl_connect('key_press_event',
self.on_key)
def on_press(self, event):
x, y = event.xdata, event.ydata
self.trajpoints.append(np.array([x, y]))
self.press = True
def on_motion(self, event):
if not self.press: return
x, y = event.xdata, event.ydata
self.ax.plot(x, y, 'r.')
self.trajpoints.append(np.array([x, y]))
self.figure.canvas.draw()
def on_release(self, event):
'on release we reset the press data'
self.press = False
traj = np.stack(self.trajpoints)
timestamps = np.linspace(0, 1, np.shape(traj)[0])
timestamps = np.expand_dims(timestamps, axis=1)
traj = np.concatenate([timestamps, traj], axis=1)
self.vmp.train(traj)
newTraj = self.vmp.roll(traj[0, 1:], traj[-1, 1:])
self.ax.cla()
self.ax.set_xlim(-5, 25)
self.ax.set_ylim(-5, 25)
self.ax.set_aspect('equal')
self.ax.plot(traj[:, 1], traj[:, 2], 'k-.')
self.ax.plot(newTraj[:, 1], newTraj[:, 2], 'g-')
self.env.plot(self.ax)
weight = self.vmp.get_flatten_weights()
starts = traj[0, 1:]
goals = traj[-1, 1:]
with open(self.bfname + "_" + "queries.csv", "a") as f:
cquery = np.expand_dims(self.current_query, axis=0)
np.savetxt(f, cquery, delimiter=',', fmt='%.3f')
with open(self.bfname + "_" + "weights.csv", "a") as f:
weight = np.expand_dims(weight, axis=1)
np.savetxt(f, weight.transpose(), delimiter=',', fmt='%.5f')
with open(self.bfname + "_" + "starts.csv", "a") as f:
starts = np.expand_dims(starts, axis=0)
np.savetxt(f, starts, delimiter=',', fmt='%.3f')
with open(self.bfname + "_" + "goals.csv", "a") as f:
goals = np.expand_dims(goals, axis=0)
np.savetxt(f, goals, delimiter=',', fmt='%.3f')
traj_fname = self.traj_dir + '/traj_' + str(self.id)
np.savetxt(traj_fname, traj, delimiter=',', fmt='%.3f')
# if self.id < self.num_data:
# self.next()
# else:
# self.disconnect()
def on_key(self, event):
if event.key == 'z':
self.next()
def next(self):
self.ax.cla()
self.ax.set_xlim(-5, 25)
self.ax.set_ylim(-5, 25)
self.ax.set_aspect('equal')
self.env = self.envs[self.id]
self.current_query = self.trqueries[self.id, :]
self.env.plot(self.ax)
self.trajpoints = []
msg = 'collect ' + str(self.id) + '-th data'
self.figure.suptitle(msg)
plt.draw()
self.id = self.id + 1
def disconnect(self):
'disconnect all the stored connection ids'
self.figure.canvas.mpl_disconnect(self.cidpress)
self.figure.canvas.mpl_disconnect(self.cidrelease)
self.figure.canvas.mpl_disconnect(self.cidmotion)
train_ws,
train_ispos,
max_epochs=5000,
is_load=False,
is_save=False)
wout, outdict = mdnmp.predict(testgoals, n_samples=n_samples)
idx = outdict['compIDs'].astype(int)
axes[expID, kid].clear()
axes[expID, kid].set_ylim([-2, 2])
axes[expID, kid].set_xlim([-2.5, 2.5])
obs.plot(axes[expID, kid])
for i in range(np.shape(vmps)[0]):
vmp.set_weights(vmps[i, :])
traj = vmp.roll(start, goals[i, :])
axes[expID, kid].plot(goals[i, 0], goals[i, 1], 'ko')
axes[expID, kid].plot(traj[:, 1],
traj[:, 2],
'k-.',
linewidth=2,
alpha=0.7)
print('the selected modes for 5 samples are {}'.format(idx))
print('probability for different modes: {}'.format(outdict['mc']))
for i in range(np.shape(wout)[0]):
cgoals = testgoals[i, :]
for j in range(n_samples):
vmp.set_weights(wout[i, j, :])
traj = vmp.roll(start, cgoals)
axes[expID, kid].plot(traj[:, 1],
gmgan = GMGAN(n_comps=2, context_dim=2, response_dim=20, nn_structure=nn_structure)
gmgan.lratio['entropy'] = 100
gmgan.gen_sup_lrate = 0.00005
gmgan.gen_adv_lrate = 0.0002
gmgan.dis_lrate = 0.0002
gmgan.entropy_ratio = 0.0
gmgan.sup_max_epoch = 5000
gmgan.create_network(num_real_data=n_data)
gmgan.init_train()
train_input = np.random.uniform(low=np.min(train_goals, axis=0), high=np.max(train_goals, axis=0),
size=(1000, np.shape(train_goals)[1]))
gmgan.train(train_context=train_input, real_context=train_goals, real_response=train_ws, max_epochs=10000, is_load=False, is_save=False)
wout, _ = gmgan.predict(testgoals)
axe.set_ylim([-2, 2])
axe.set_xlim([-2.5, 2.5])
obs.plot(axe)
for k in range(np.shape(wout)[0]):
cgoals = testgoals[k, :]
vmp.set_weights(wout[k,0,:])
traj = vmp.roll(start, cgoals)
axe.plot(traj[:, 1], traj[:, 2], '-', color='r', linewidth=2)
axe.plot(cgoals[0], cgoals[1], 'bo')
plt.draw()
plt.show()
def evaluate_docking_for_all_models(mpweights,
mfs,
testqueries,
teststarts,
testgoals,
knum=10,
ndraw=4):
ndata = np.shape(mpweights)[0]
vmp = VMP(2, kernel_num=knum)
obsExp = ObsExp(exp_name="Docking")
envs = obsExp.get_envs(testqueries)
fig, axes = plt.subplots(nrows=ndraw, ncols=ndraw)
colors = [(232 / 255, 82 / 255, 88 / 255, 1),
(1 / 255, 129 / 255, 1 / 255, 1),
(4 / 255, 129 / 255, 205 / 255, 1)
] #cmap(np.linspace(0,1,nmodel))
k = 0
for i in range(ndata):
ws = mpweights[i, :, :]
iscollision = False
for j in range(np.shape(ws)[1]):
start = teststarts[i, :]
goal = testgoals[i, :]
env = envs[i]
mf = mfs[i, :]
vmp.set_weights(ws[:, np.argmax(mf)])
traj = vmp.roll(start, goal)
if np.linalg.norm(traj[0, 1:] -
start) > 1 or np.linalg.norm(traj[-1, 1:] -
goal) > 1:
iscollision = True
continue
if np.linalg.norm(testqueries[i, 0:2] -
start) > 1 or np.linalg.norm(
testqueries[i, 3:5] - goal) > 1:
iscollision = True
continue
for j in range(np.shape(traj)[0]):
iscollision = env.isCollision([traj[j, 1], traj[j, 2]])
if iscollision:
break
if not iscollision and k < ndraw * ndraw:
ax = int(np.floor(k / ndraw))
ay = int(k % ndraw)
axes[ax, ay].clear()
axes[ax, ay].set_ylim([-5, 25])
axes[ax, ay].set_xlim([-5, 25])
axes[ax, ay].set_yticklabels([])
axes[ax, ay].set_xticklabels([])
axes[ax, ay].set_aspect('equal')
axes[ax, ay].tick_params(length=0)
env.plot(axes[ax, ay])
axes[ax, ay].plot(traj[0, 1], traj[0, 2], 'k.')
axes[ax, ay].plot(traj[-1, 1], traj[-1, 2], 'k.')
for mid in range(len(mf)):
vmp.set_weights(ws[:, mid])
traj = vmp.roll(start, goal)
if mid == np.argmax(mf):
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'-',
color=colors[mid],
linewidth=2)
else:
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'-.',
color=colors[mid],
linewidth=2)
k = k + 1
plt.subplots_adjust(hspace=0.05, wspace=0.05)
plt.draw()
plt.pause(0.001)
plt.savefig('./docking_img.png')
def evaluate_docking(mpweights,
testqueries,
teststarts,
testgoals,
knum=10,
mua=None,
isdraw=False,
ndraw=3,
onlySuccess=True):
# mpweights: N x S x dim, N: number of experiments, S: number of samples, dim: dimension of MP
# testqueries: N x qdim, N: number of experiments, qdim: dimension of queries
# teststarts: N x 2
# testgoals: N x 2
success_num = 0
ndata = np.shape(mpweights)[0]
vmp = VMP(2, kernel_num=knum)
obsExp = ObsExp(exp_name="Docking")
envs = obsExp.get_envs(testqueries)
k = 0
if isdraw:
if not onlySuccess:
ids = np.random.choice(ndata, ndraw * ndraw, replace=False)
else:
ids = np.linspace(0, ndata, ndata)
fig, axes = plt.subplots(nrows=ndraw, ncols=ndraw)
successId = []
colors = [(232 / 255, 82 / 255, 88 / 255, 0.5),
(1 / 255, 129 / 255, 1 / 255, 0.5),
(4 / 255, 129 / 255, 205 / 255, 0.5)]
for i in range(ndata):
w = mpweights[i, :]
for j in range(np.shape(w)[0]):
vmp.set_weights(w[j, :])
start = teststarts[i, :]
goal = testgoals[i, :]
env = envs[i]
traj = vmp.roll(start, goal)
iscollision = False
if np.linalg.norm(traj[0, 1:] -
start) > 1 or np.linalg.norm(traj[-1, 1:] -
goal) > 1:
iscollision = True
continue
for j in range(np.shape(traj)[0]):
iscollision = env.isCollision([traj[j, 1], traj[j, 2]])
if iscollision:
break
if not iscollision:
success_num = success_num + 1
successId.append(i)
break
if isdraw:
if k < ndraw * ndraw:
if not onlySuccess or not iscollision:
ax = int(np.floor(k / ndraw))
ay = int(k % ndraw)
axes[ax, ay].clear()
axes[ax, ay].set_ylim([-5, 25])
axes[ax, ay].set_xlim([-5, 25])
axes[ax, ay].set_yticklabels([])
axes[ax, ay].set_xticklabels([])
axes[ax, ay].set_aspect('equal')
axes[ax, ay].tick_params(length=0)
# axes[ax, ay].axhline(linewidth=2)
# axes[ax, ay].axvline(linewidth=2)
env.plot(axes[ax, ay])
axes[ax, ay].plot(start[0], start[1], 'r.')
axes[ax, ay].plot(goal[0], goal[1], 'b.')
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'k-',
linewidth=2)
if mua is not None:
for mid in range(np.shape(mua)[2]):
vmp.set_weights(mua[i, :, mid])
traj = vmp.roll(start, goal)
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'-.',
color=colors[mid],
linewidth=2)
k = k + 1
print('success_num: %1d, ndata: %1d, success_rate: %.3f' %
(success_num, ndata, success_num / ndata))
if isdraw:
plt.subplots_adjust(hspace=0.05, wspace=0.05)
plt.draw()
plt.pause(0.001)
plt.savefig('./docking.png', format='png', dpi=1200)
return success_num / ndata, successId
def evaluate_circular_obstacle_avoidance(mpweights,
testqueries,
teststarts,
testgoals,
knum=10,
isdraw=False,
ndraw=4,
onlySuccess=False):
# mpweights: N x dim, N: number of experiments, dim: dimension of MP
# testqueries: N x qdim, N: number of experiments, qdim: dimension of queries
# teststarts: N x 2
# testgoals: N x 2
success_num = 0
ndata = np.shape(mpweights)[0]
vmp = VMP(2, kernel_num=knum)
obsExp = ObsExp(exp_name="goAroundObsR3V2")
envs = obsExp.get_envs(testqueries)
k = 0
if isdraw:
if not onlySuccess:
ids = np.random.choice(ndata, ndraw * ndraw, replace=False)
else:
ids = np.linspace(0, ndata, ndata)
fig, axes = plt.subplots(nrows=ndraw, ncols=ndraw)
for i in range(ndata):
w = mpweights[i, :]
vmp.set_weights(w)
start = teststarts[i, :]
goal = testgoals[i, :]
env = envs[i]
traj = vmp.roll(start, goal)
iscollision = False
for j in range(np.shape(traj)[0]):
iscollision = env.isCollision([traj[j, 1], traj[j, 2]])
if iscollision:
break
if not iscollision:
success_num = success_num + 1
if isdraw:
if i in ids:
if not onlySuccess or not iscollision:
ax = int(np.floor(k / ndraw))
ay = int(k % ndraw)
axes[ax, ay].clear()
axes[ax, ay].set_ylim([-10, 10])
axes[ax, ay].set_xlim([-10, 10])
axes[ax, ay].set_yticklabels([])
axes[ax, ay].set_xticklabels([])
env.plot(axes[ax, ay])
k = k + 1
if iscollision:
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'-.',
color='r')
else:
axes[ax, ay].plot(traj[:, 1],
traj[:, 2],
'-',
color='b')
axes[ax, ay].plot(start[0], start[1], 'ro')
axes[ax, ay].plot(goal[0], goal[1], 'bo')
print('success_num: {}, ndata: {}'.format(success_num, ndata))
if isdraw:
plt.show()
return success_num / ndata