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:
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))
gui.update()
if output:
gui.screenshot('mass_spring/{}/{:04d}.png'.format(output, t))
loss[None] = 0
compute_loss(steps - 1)
def visualize(s, folder):
canvas.clear(0xFFFFFF)
vec = tc.vec
for i in range(n_particles):
color = 0x111111
aid = actuator_id[i]
if aid != -1:
act = actuation[s - 1, aid]
color = rgb_to_hex((0.5 - act, 0.5 - abs(act), 0.5 + act))
canvas.circle(vec(x[s, i][0], x[s, i][1])).radius(2).color(color).finish()
canvas.path(tc.vec(0.05, 0.02), tc.vec(0.95, 0.02)).radius(3).color(0x0).finish()
gui.update()
os.makedirs(folder, exist_ok=True)
gui.screenshot('{}/{:04d}.png'.format(folder, s))
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):
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:
canvas.clear(0xFFFFFF)
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):
canvas.path(tc.vec(*points[k]), tc.vec(
*points[(k + 1) % 4])).color(0x0).radius(2).finish()
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 = rgb_to_hex((0.5 + a, 0.5 - abs(a), 0.5 - a))
if spring_length[i] == -1:
canvas.path(
tc.vec(*pt1),
tc.vec(*pt2)).color(0x000000).radius(9).finish()
canvas.path(tc.vec(*pt1),
tc.vec(*pt2)).color(color).radius(7).finish()
else:
canvas.path(
tc.vec(*pt1),
tc.vec(*pt2)).color(0x000000).radius(7).finish()
canvas.path(tc.vec(*pt1),
tc.vec(*pt2)).color(color).radius(5).finish()
canvas.path(tc.vec(0.05, ground_height - 5e-3),
tc.vec(0.95, ground_height -
5e-3)).color(0x0).radius(5).finish()
gui.update()
if output:
gui.screenshot('rigid_body/{}/{:04d}.png'.format(output, t))
loss[None] = 0
compute_loss(steps - 1)
