def draw_gui():
X = position.to_numpy()
Y = rest_length.to_numpy()
Z = adj_ptr.to_numpy()
# gui.circles(X[:num_particles[None]], color=0xffaa77, radius=1)
gui.line(begin=(left_x, 0.0), end=(left_x, 1.0), color=0x0, radius=1)
for i in range(num_particles[None]):
for ptr_j, j in enumerate(Z[i]):
if j >= 0 and i > j:
dist = ((X[i][0] - X[j][0])**2 + (X[i][1] - X[j][1])**2)**0.5
ratio = (dist - Y[i][ptr_j]) / Y[i][ptr_j]
ratio *= ratio_mul
ratio = max(-1.0, ratio)
ratio = min(1.0, ratio)
if ratio < 0:
color = ti.rgb_to_hex((1.0, 1.0 + ratio, 1.0 + ratio))
else:
color = ti.rgb_to_hex((1.0 - ratio, 1.0 - ratio, 1.0))
gui.line(begin=X[i], end=X[j], radius=1.5, color=color)
inv_freq = vib_dis[0] * (1.0 / 24.0 / steps)
freq = 1.0 / inv_freq
gui.text(content=f'C: clear all', pos=(left_x + 0.01, 0.99), color=0x0)
gui.text(content=f'Space: pause', pos=(left_x + 0.01, 0.96), color=0x0)
gui.text(content=f'S: Spring stiffness {spring_stiffness[None]:.1f}',
pos=(left_x + 0.01, 0.93),
color=0x0)
gui.text(content=f'F: Freq {freq:.4f}Hz T={inv_freq:.4f}s',
pos=(left_x + 0.01, 0.90),
color=0x0)
def paint_phi(gui):
pos_ = pos.to_numpy()
phi_ = phi.to_numpy()
f2v_ = f2v.to_numpy()
a, b, c = pos_[f2v_[:, 0]], pos_[f2v_[:, 1]], pos_[f2v_[:, 2]]
k = phi_ * (10 / E)
gb = (1 - k) * 0.5
gui.triangles(a, b, c, color=ti.rgb_to_hex([k + gb, gb, gb]))
def visualize(s, folder):
aid = actuator_id.to_numpy()
for i in range(n_particles):
color = 0x111111
if aid[i] != -1:
act = actuation[s - 1, aid[i]]
color = ti.rgb_to_hex((0.5 - act, 0.5 - abs(act), 0.5 + act))
gui.circle([x[s, i][0], x[s, i][1]], radius=1.5, color=color)
gui.line((0.05, 0.02), (0.95, 0.02), radius=3, color=0x0)
os.makedirs(folder, exist_ok=True)
gui.show(f'{folder}/{s:04d}.png')
def paint_phi(gui):
pos_np = pos.to_numpy()
phi_np = phi.to_numpy()
f2v_np = f2v.to_numpy()
a, b, c = pos_np[f2v_np[:, 0]], pos_np[f2v_np[:, 1]], pos_np[f2v_np[:, 2]]
k = phi_np * (8000 / E)
gb = (1 - k) * 0.7
# print("gb:", gb[0])
# print("phi_np", phi_np[0])
# print("k", k[0])
gui.triangles(a, b, c, color=ti.rgb_to_hex([k + gb, gb, gb]))
gui.lines(a, b, color=0xffffff, radius=0.5)
gui.lines(b, c, color=0xffffff, radius=0.5)
gui.lines(c, a, color=0xffffff, radius=0.5)
def forward(output="run", visualize=True):
interval = vis_interval
if output:
interval = output_vis_interval
os.makedirs('mass_spring/{}/'.format(output), exist_ok=True)
total_steps = steps if not output else steps * 2
for t in range(1, total_steps):
print(x.to_numpy())
compute_center(t - 1)
apply_spring_force(t - 1)
advance_toi(t)
if (t + 1) % interval == 0 and visualize:
gui.clear()
gui.line((0, ground_height), (1, ground_height),
color=0x0,
radius=3)
def circle(x, y, color):
gui.circle((x, y), ti.rgb_to_hex(color), 7)
for i in range(n_springs):
def get_pt(x):
return (x[0], x[1])
a = 0 * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
gui.line(get_pt(x[t, spring_anchor_a[i]]),
get_pt(x[t, spring_anchor_b[i]]),
color=c,
radius=r)
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(x[t, i][0], x[t, i][1], color)
if output:
gui.show('mass_spring/{}/{:04d}.png'.format(output, t))
else:
gui.show()
def visualize(s, folder):
aid = actuator_id.to_numpy()
colors = np.empty(shape=n_particles, dtype=np.uint32)
particles = x.to_numpy()[s]
for i in range(n_particles):
color = 0x111111
if aid[i] != -1:
act = actuation[s - 1, aid[i]]
color = ti.rgb_to_hex((0.5 - act, 0.5 - abs(act), 0.5 + act))
colors[i] = color
gui.circles(pos=particles, color=colors, radius=1.5)
gui.line((0.05, 0.02), (0.95, 0.02), radius=3, color=0x0)
os.makedirs(folder, exist_ok=True)
gui.show(f'{folder}/{s:04d}.png')
def animate(xs, acts, ground_height: float, outputdir=None, head_id=0):
"""Animate a the policy controlling the robot.
* `total_steps` controls the number of time steps to animate for. This is
necessary since the final animation is run for more time steps than in
training.
* `outputdir` controls whether or not frames of the animation are screenshotted
and dumped to a results directory.
"""
assert len(xs) == len(acts)
gui = ti.GUI("Mass Spring Robot", (512, 512),
background_color=0xFFFFFF,
show_gui=False)
for t in range(len(xs)):
gui.line(begin=(0, ground_height),
end=(1, ground_height),
color=0x0,
radius=3)
def circle(x, y, color):
gui.circle((x, y), ti.rgb_to_hex(color), 7)
for i in range(n_springs):
a = acts[t][i] * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
gui.line(begin=tuple(xs[t][spring_anchor_a[i], :]),
end=tuple(xs[t][spring_anchor_b[i], :]),
radius=r,
color=c)
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(xs[t][i][0], xs[t][i][1], color)
if outputdir is not None:
gui.show(os.path.join(outputdir, "{:04d}.png".format(t + 1)))
else:
gui.show()
def animate(total_steps: int, output=None):
"""Animate a the policy controlling the robot.
* `total_steps` controls the number of time steps to animate for. This is
necessary since the final animation is run for more time steps than in
training.
* `output` controls whether or not frames of the animation are screenshotted
and dumped to a results directory.
"""
if output:
os.makedirs("{}/{}/".format(RESULTS_DIR, output))
for t in range(1, total_steps):
canvas.clear(0xFFFFFF)
canvas.path(ti.vec(0, ground_height),
ti.vec(1, ground_height)).color(0x0).radius(3).finish()
def circle(x, y, color):
canvas.circle(ti.vec(x, y)).color(
ti.rgb_to_hex(color)).radius(7).finish()
for i in range(n_springs):
def get_pt(x):
return ti.vec(x[0], x[1])
a = act[t - 1, i] * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
canvas.path(
get_pt(x[t, spring_anchor_a[i]]),
get_pt(x[t, spring_anchor_b[i]])).color(c).radius(r).finish()
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(x[t, i][0], x[t, i][1], color)
gui.update()
if output:
gui.screenshot("{}/{}/{:04d}.png".format(RESULTS_DIR, output, t))
def circle(x, y, color):
canvas.circle(tc.vec(x, y)).color(
rgb_to_hex(color)).radius(7).finish()
def forward(output=None, visualize=True):
if random.random() > 0.5:
goal[None] = [0.9, 0.2]
else:
goal[None] = [0.1, 0.2]
goal[None] = [0.9, 0.2]
interval = vis_interval
if output:
interval = output_vis_interval
os.makedirs('mass_spring/{}/'.format(output), exist_ok=True)
total_steps = steps if not output else steps * 2
pool = [(random.random() - 0.5) * 2 for _ in range(100)]
for i in range(total_steps):
if output:
target_v[i][0] = (i // 300) % 2 * 2 - 1
else:
target_v[i][0] = pool[i // 300]
for t in range(1, total_steps):
compute_center(t - 1)
nn1(t - 1)
nn2(t - 1)
apply_spring_force(t - 1)
if use_toi:
advance_toi(t)
else:
advance_no_toi(t)
compute_loss(t)
if (t + 1) % interval == 0 and visualize:
canvas.clear(0xFFFFFF)
canvas.path(tc.vec(0, ground_height),
tc.vec(1,
ground_height)).color(0x0).radius(3).finish()
def circle(x, y, color):
canvas.circle(tc.vec(x, y)).color(
rgb_to_hex(color)).radius(7).finish()
for i in range(n_springs):
def get_pt(x):
return tc.vec(x[0], x[1])
a = act[t - 1, i] * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
canvas.path(
get_pt(x[t, spring_anchor_a[i]]),
get_pt(x[t,
spring_anchor_b[i]])).color(c).radius(r).finish()
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(x[t, i][0], x[t, i][1], color)
# circle(goal[None][0], goal[None][1], (0.6, 0.2, 0.2))
if target_v[t][0] > 0:
circle(0.5, 0.5, (1, 0, 0))
circle(0.6, 0.5, (1, 0, 0))
else:
circle(0.5, 0.5, (0, 0, 1))
circle(0.4, 0.5, (0, 0, 1))
gui.update()
if output:
gui.screenshot('mass_spring/{}/{:04d}.png'.format(output, t))
def main():
tc.set_gdb_trigger()
# initialization
scene = Scene()
# fish(scene)
robot(scene)
# scene.add_rect(0.4, 0.4, 0.2, 0.1, 0.3, 0.1, -1, 1)
scene.finalize()
for i in range(n_actuators):
for j in range(n_sin_waves):
weights[i, j] = np.random.randn() * 0.01
for i in range(scene.n_particles):
x[0, i] = scene.x[i]
F[0, i] = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
actuator_id[i] = scene.actuator_id[i]
particle_type[i] = scene.particle_type[i]
losses = []
for iter in range(100):
t = time.time()
ti.clear_all_gradients()
l = forward()
losses.append(l)
loss.grad[None] = 1
backward()
per_iter_time = time.time() - t
print('i=', iter, 'loss=', l, F' per iter {per_iter_time:.2f}s')
learning_rate = 30
for i in range(n_actuators):
for j in range(n_sin_waves):
weights[i, j] -= learning_rate * weights.grad[i, j]
bias[i] -= learning_rate * bias.grad[i]
if iter % 20 == 19:
print('Writing particle data to disk...')
print('(Please be patient)...')
# visualize
forward()
x_ = x.to_numpy()
v_ = v.to_numpy()
particle_type_ = particle_type.to_numpy()
actuation_ = actuation.to_numpy()
actuator_id_ = actuator_id.to_numpy()
folder = 'mpm3d/iter{:04d}/'.format(iter)
os.makedirs(folder, exist_ok=True)
for s in range(7, steps, 2):
xs, ys, zs = [], [], []
us, vs, ws = [], [], []
cs = []
for i in range(n_particles):
xs.append(x_[s, i][0])
ys.append(x_[s, i][1])
zs.append(x_[s, i][2])
us.append(v_[s, i][0])
vs.append(v_[s, i][1])
ws.append(v_[s, i][2])
if particle_type_[i] == 0:
# fluid
r = 0.3
g = 0.3
b = 1.0
else:
# neohookean
if actuator_id_[i] != -1:
# actuated
act = actuation_[s, actuator_id_[i]] * 0.5
r = 0.5 - act
g = 0.5 - abs(act)
b = 0.5 + act
else:
r, g, b = 0.4, 0.4, 0.4
cs.append(ti.rgb_to_hex((r, g, b)))
data = np.array(xs + ys + zs + us + vs + ws + cs,
dtype=np.float32)
fn = '{}/{:04}.bin'.format(folder, s)
data.tofile(open(fn, 'wb'))
print('.', end='')
print()
plt.title("Optimization of Initial Velocity")
plt.ylabel("Loss")
plt.xlabel("Gradient Descent Iterations")
plt.plot(losses)
plt.show()
def forward(output=None, visualize=True):
initialize_properties()
interval = vis_interval
total_steps = steps
if output:
print(output)
interval = output_vis_interval
os.makedirs('rigid_body/{}/'.format(output), exist_ok=True)
total_steps *= 2
goal[None] = [0.9, 0.15]
for t in range(1, total_steps):
nn1(t - 1)
nn2(t - 1)
collide(t - 1)
apply_spring_force(t - 1)
if use_toi:
advance_toi(t)
else:
advance_no_toi(t)
if (t + 1) % interval == 0 and visualize:
for i in range(n_objects):
points = []
for k in range(4):
offset_scale = [[-1, -1], [1, -1], [1, 1], [-1, 1]][k]
rot = rotation[t, i]
rot_matrix = np.array([[math.cos(rot), -math.sin(rot)],
[math.sin(rot),
math.cos(rot)]])
pos = np.array([x[t, i][0], x[t, i][1]
]) + offset_scale * rot_matrix @ np.array(
[halfsize[i][0], halfsize[i][1]])
points.append((pos[0], pos[1]))
for k in range(4):
gui.line(points[k],
points[(k + 1) % 4],
color=0x0,
radius=2)
for i in range(n_springs):
def get_world_loc(i, offset):
rot = rotation[t, i]
rot_matrix = np.array([[math.cos(rot), -math.sin(rot)],
[math.sin(rot),
math.cos(rot)]])
pos = np.array([[x[t, i][0]], [
x[t, i][1]
]]) + rot_matrix @ np.array([[offset[0]], [offset[1]]])
return pos
pt1 = get_world_loc(spring_anchor_a[i], spring_offset_a[i])
pt2 = get_world_loc(spring_anchor_b[i], spring_offset_b[i])
color = 0xFF2233
if spring_actuation[i] != 0 and spring_length[i] != -1:
a = actuation[t - 1, i] * 0.5
color = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
if spring_length[i] == -1:
gui.line(pt1, pt2, color=0x000000, radius=9)
gui.line(pt1, pt2, color=color, radius=7)
else:
gui.line(pt1, pt2, color=0x000000, radius=7)
gui.line(pt1, pt2, color=color, radius=5)
gui.line((0.05, ground_height - 5e-3),
(0.95, ground_height - 5e-3),
color=0x0,
radius=5)
file = None
if output:
file = f'rigid_body/{output}/{t:04d}.png'
gui.show(file=file)
loss[None] = 0
compute_loss(steps - 1)
def circle(x, y, color):
gui.circle((x, y), ti.rgb_to_hex(color), 7)
def forward(output=None, visualize=True):
if random.random() > 0.5:
goal[None] = [0.9, 0.2]
else:
goal[None] = [0.1, 0.2]
goal[None] = [0.9, 0.2]
interval = vis_interval
if output:
interval = output_vis_interval
os.makedirs('mass_spring/{}/'.format(output), exist_ok=True)
total_steps = steps if not output else steps * 2
for t in range(1, total_steps):
compute_center(t - 1)
nn1(t - 1)
nn2(t - 1)
apply_spring_force(t - 1)
if use_toi:
advance_toi(t)
else:
advance_no_toi(t)
if (t + 1) % interval == 0 and visualize:
gui.clear()
gui.line((0, ground_height), (1, ground_height), 0x0, 3)
def circle(x, y, color):
gui.circle((x, y), ti.rgb_to_hex(color), 7)
for i in range(n_springs):
def get_pt(x):
return (x[0], x[1])
a = act[t - 1, i] * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
gui.line(get_pt(x[t, spring_anchor_a[i]]),
get_pt(x[t, spring_anchor_b[i]]), c, r)
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(x[t, i][0], x[t, i][1], color)
# circle(goal[None][0], goal[None][1], (0.6, 0.2, 0.2))
if output:
gui.show('mass_spring/{}/{:04d}.png'.format(output, t))
else:
gui.show()
loss[None] = 0
compute_loss(steps - 1)
import taichi as ti
ti.init(default_fp=ti.f32, arch=ti.cpu)
rgb = (0.4, 0.8, 1.0)
clr = ti.rgb_to_hex(rgb) # 0x66ccff
gui = ti.GUI("velocity plot", (512, 512), background_color=0x000fff)
gui.circle(pos=(0, 0), color=0xffffff, radius=100)
gui.text(content="Ni ma ge cou bi!", pos=(0.5, 0.1), font_size=20, color=clr)
gui.text(content="Ni ma ge cou bi!", pos=(0.5, 0.5), font_size=20, color=clr)
gui.text(content="Ni ma ge cou bi!", pos=(0.5, 0.8), font_size=20, color=clr)
gui.show("sample.png")
def forward(output=None, visualize=True):
if random.random() > 0.5:
goal[None] = [0.9, 0.2]
else:
goal[None] = [0.1, 0.2]
goal[None] = [0.9, 0.2]
interval = vis_interval
if output:
interval = output_vis_interval
os.makedirs('mass_spring/{}/'.format(output), exist_ok=True)
total_steps = steps if not output else steps * 2
pool = [(random.random() - 0.5) * 2 for _ in range(100)]
for i in range(total_steps):
if output:
target_v[i][0] = ((i // turn_period) % 2 * 2 - 1) * 1
else:
target_v[i][0] = (pool[i // turn_period] * 2) * 1
if target_v[i][0] < 0:
target_v[i][0] = -target_v[i][0]
target_v[i][0] = 0.2
if output:
target_h[None] = 0.5
else:
target_h[None] = 0.4 + random.random() * 0.2
for t in range(1, total_steps):
compute_center(t - 1)
nn1(t - 1)
nn2(t - 1)
apply_spring_force(t - 1)
advance_toi(t)
if duplicate_v > 0:
compute_loss(t)
if duplicate_h > 0 and t % cycle_period == cycle_period // 2:
compute_loss_h(t - 1)
# output_target.append(target_v[t][0])
# output_sim.append(v[t, head_id][0])
output_target.append(target_h[None])
output_sim.append(x[t, head_id][1])
if (t + 1) % interval == 0 and visualize:
gui.clear()
gui.line((0, ground_height * 0.5), (1, ground_height * 0.5),
color=0x0,
radius=3)
def circle(x, y, color):
gui.circle((x, y), ti.rgb_to_hex(color), 3)
for i in range(n_springs):
def get_pt(x):
return (x[0] * 0.5 + 0.2, x[1] * 0.5)
a = act[t - 1, i] * 0.5
r = 2
if spring_actuation[i] == 0:
a = 0
c = 0x222222
else:
r = 4
c = ti.rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
gui.line(get_pt(x[t, spring_anchor_a[i]]),
get_pt(x[t, spring_anchor_b[i]]),
color=c,
radius=r / 2)
for i in range(n_objects):
color = (0.4, 0.6, 0.6)
if i == head_id:
color = (0.8, 0.2, 0.3)
circle(x[t, i][0] * 0.5 + 0.2, x[t, i][1] * 0.5, color)
# circle(goal[None][0], goal[None][1], (0.6, 0.2, 0.2))
if target_v[t][0] > 0:
circle(0.5, 0.5, (1, 0, 0))
circle(0.6, 0.5, (1, 0, 0))
else:
circle(0.5, 0.5, (0, 0, 1))
circle(0.4, 0.5, (0, 0, 1))
if output:
gui.show('mass_spring/{}/{:04d}.png'.format(output, t))
else:
gui.show()
compute_loss_x(total_steps - 1)
output_loss.append(loss[None])
fig = plt.figure()
temp_loss = gaussian_filter(output_loss, 10)
plt.plot(temp_loss)
fig.savefig('plots/' + str(forward.cnt) + '.png', dpi=fig.dpi)
plt.close(fig)