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)