def testChain(self):
bodies = []
joints = []
# make chain of rectangles
r = Rect([300, 50], [20, 60])
bodies.append(r)
joints.append(XConstraint(r))
joints.append(YConstraint(r))
for i in range(1, 10):
r = Rect([300, 50 + 50 * i], [20, 60])
bodies.append(r)
joints.append(Joint(bodies[-1], bodies[-2], [300, 25 + 50 * i]))
bodies[-1].add_no_collision(bodies[-2])
bodies[-1].add_force(ExternalForce(down_force, multiplier=100))
# make projectile
c = Circle([50, 500], 20, restitution=1)
bodies.append(c)
c.add_force(ExternalForce(hor_impulse, multiplier=1000))
recorder = None
# recorder = Recorder(DT, self.screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=TIME, screen=self.screen, recorder=recorder)
def fixed_joint_demo(screen):
bodies = []
joints = []
restitution = 0.5
fric_coeff = 0.15
r = Rect([120, 100], [60, 60],
restitution=restitution, fric_coeff=fric_coeff)
bodies.append(r)
r.add_force(ExternalForce(gravity, multiplier=100))
r2 = Rect([160, 100], [60, 60],
restitution=restitution, fric_coeff=fric_coeff)
bodies.append(r2)
joints += [FixedJoint(r, r2)]
r2.add_no_collision(r)
r2.add_force(ExternalForce(gravity, multiplier=100))
inclination = math.pi / 32
r = Rect([inclination, 500, 500], [900, 10],
restitution=restitution, fric_coeff=fric_coeff)
bodies.append(r)
joints.append(TotalConstraint(r))
recorder = None
# recorder = Recorder(DT, screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=TIME, screen=screen, recorder=recorder)
def chain_demo(screen):
bodies = []
joints = []
restitution = 0.9
# make chain of rectangles
r = Rect([300, 50], [20, 60], restitution=restitution)
bodies.append(r)
joints.append(XConstraint(r))
joints.append(YConstraint(r))
for i in range(1, 10):
r = Rect([300, 50 + 50 * i], [20, 60], restitution=restitution)
bodies.append(r)
joints.append(Joint(bodies[-1], bodies[-2], [300, 25 + 50 * i]))
bodies[-1].add_no_contact(bodies[-2])
bodies[-1].add_force(Gravity(g=100))
# make projectile
c = Circle([50, 500], 20, restitution=restitution)
bodies.append(c)
c.add_force(ExternalForce(hor_impulse, multiplier=2000))
clock = Circle([975, 575], 20, vel=[1, 0, 0])
bodies.append(clock)
recorder = None
# recorder = Recorder(DT, screen)
world = World(bodies, joints, dt=DT, post_stab=True)
run_world(world, run_time=TIME * 2, screen=screen, recorder=recorder)
def testFric(self):
bodies = []
joints = []
def timed_force(t):
if 1 < t < 2:
return ExternalForce.RIGHT
else:
return ExternalForce.ZEROS
r = Rect([400, 400], [900, 10])
bodies.append(r)
r.add_force(ExternalForce(timed_force, multiplier=100))
r.add_force(ExternalForce(gravity, multiplier=100))
c = Circle([200, 364], 30)
bodies.append(c)
c.add_force(ExternalForce(gravity, multiplier=100))
c = Circle([50, 436], 30)
bodies.append(c)
joints.append(XConstraint(c))
joints.append(YConstraint(c))
c = Circle([800, 436], 30)
bodies.append(c)
joints.append(XConstraint(c))
joints.append(YConstraint(c))
recorder = None
# recorder = Recorder(DT, self.screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=10, screen=self.screen, recorder=recorder)
def testSlide(self):
bodies = []
joints = []
restitution = 0.5
fric_coeff = 0.2
clock = Circle([975, 575], 20, vel=[1, 0, 0])
bodies.append(clock)
p1 = Hull([500, 300], [[450, 5], [-450, 5], [-450, -5], [450, -5]],
restitution=restitution, fric_coeff=fric_coeff)
p1.rotate_verts(get_tensor(math.pi / 32))
bodies.append(p1)
joints.append(TotalConstraint(p1))
# Rectangle
# p2 = Hull([100, 100], [[30, 30], [-30, 30], [-30, -30], [30, -30]],
# restitution=restitution, fric_coeff=fric_coeff)
# Pentagon
p2 = Hull([100, 100], [[50, 0], [30, 50], [-30, 30], [-50, -30], [0, -50]],
restitution=restitution, fric_coeff=fric_coeff)
# Hexagon
p2 = Hull([100, 100], [[50, 0], [30, 50], [-30, 30], [-50, -30], [0, -50], [30, -30]],
restitution=restitution, fric_coeff=fric_coeff)
bodies.append(p2)
p2.add_force(ExternalForce(down_force, multiplier=100))
# p2.add_force(ExternalForce(hor_impulse, multiplier=-100))
recorder = None
# recorder = Recorder(DT, self.screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=TIME, screen=self.screen, recorder=recorder)
def testFric(self):
restitution = 0.75
fric_coeff = 1
bodies = []
joints = []
def timed_force(t):
if 1 < t < 2:
return ExternalForce.RIGHT
else:
return ExternalForce.ZEROS
r = Rect([400, 400], [900, 10],
restitution=restitution,
fric_coeff=fric_coeff)
bodies.append(r)
r.add_force(ExternalForce(timed_force, multiplier=100))
r.add_force(ExternalForce(down_force, multiplier=100))
c = Circle([200, 364],
30,
restitution=restitution,
fric_coeff=fric_coeff)
bodies.append(c)
c.add_force(ExternalForce(down_force, multiplier=100))
c = Circle([50, 436],
30,
restitution=restitution,
fric_coeff=fric_coeff)
bodies.append(c)
joints.append(XConstraint(c))
joints.append(YConstraint(c))
c = Circle([800, 436],
30,
restitution=restitution,
fric_coeff=fric_coeff)
bodies.append(c)
joints.append(XConstraint(c))
joints.append(YConstraint(c))
clock = Circle([975, 575], 20, vel=[1, 0, 0])
bodies.append(clock)
recorder = None
# recorder = Recorder(DT, screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=10, screen=self.screen, recorder=recorder)
def testDemo(self):
bodies = []
joints = []
# Ball hitting object constrained by 1 joint
for i in range(1, 3):
c = Circle([150, 150 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(vert_impulse, multiplier=500))
bodies.append(c)
joints.append(Joint(bodies[-1], None, [140, 220]))
# Ball bouncing on body fixed in place
for i in range(1, 3):
c = Circle([300 + 1 * (i - 1), 150 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(down_force, multiplier=100))
bodies.append(c)
joints.append(TotalConstraint(bodies[-1]))
# 2 free ball collision angled
for i in range(1, 3):
c = Circle([225 - 10 * (i - 1), 300 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(down_force, multiplier=100))
bodies.append(c)
# 2 free ball collision straight
for i in range(1, 3):
c = Circle([375, 300 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(vert_impulse, multiplier=500))
bodies.append(c)
r = Rect([300, 500], [40, 40])
r.add_force(ExternalForce(down_force, multiplier=-100))
r.v[0] = -1
bodies.append(r)
r = Rect([300, 50], [40, 40])
r.add_force(ExternalForce(down_force, multiplier=100))
r.v[0] = -1
for b in bodies:
b.add_no_collision(r)
# bodies.append(r)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=10, screen=self.screen)
def debug_demo(screen):
bodies = []
joints = []
# Ball hitting object constrained by 1 joint
for i in range(1, 3):
c = Circle([150, 150 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(vert_impulse, multiplier=500))
bodies.append(c)
joints.append(Joint(bodies[-1], None, [140, 220]))
# Ball bouncing on body fixed in place
for i in range(1, 3):
c = Circle([300 + 1 * (i - 1), 150 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(Gravity(g=100))
# else:
# c.add_force(ExternalForce(neg_gravity, multiplier=100))
bodies.append(c)
joints.append(TotalConstraint(bodies[-1]))
# 2 free ball collision angled
for i in range(1, 3):
c = Circle([225 - 10 * (i - 1), 300 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(Gravity(g=100))
bodies.append(c)
# 2 free ball collision straight
for i in range(1, 3):
c = Circle([375, 300 + 80 * (i - 1)], 20)
if i == 1:
c.add_force(ExternalForce(vert_impulse, multiplier=500))
bodies.append(c)
r = Rect([300, 500], [40, 40], vel=[-1, 0, 0])
r.add_force(Gravity(g=-100))
bodies.append(r)
clock = Circle([975, 575], 20, vel=[1, 0, 0])
bodies.append(clock)
world = World(bodies, joints, dt=DT, post_stab=True)
run_world(world, run_time=10, screen=screen)
def testRender(self):
p1 = Hull([700, 200], [[math.sqrt(2) * 10 + 10, 0 + 10], [0 + 10, math.sqrt(2) * 10 + 10],
[math.sqrt(2) * -10 + 10, 0 + 10], [0 + 10, math.sqrt(2) * -10 + 10]])
p2 = Hull([300, 200], [[60, 0], [60, 60], [0, 60]])
p2.add_force(ExternalForce(hor_impulse, multiplier=1000))
p2.add_force(ExternalForce(rot_impulse, multiplier=1000))
p3 = Hull([700, 400], [[60, 0], [60, 60], [0, 0]])
p4 = Hull([300, 420], [[60, 0], [0, 60], [0, 0]])
p4.add_force(ExternalForce(hor_impulse, multiplier=1000))
p4.add_force(ExternalForce(rot_impulse, multiplier=1000))
p5 = Hull([500, 300], [[50, 0], [30, 60], [-30, 30], [-60, -30], [0, -60]])
p5.add_force(ExternalForce(hor_impulse, multiplier=1000))
p5.add_force(ExternalForce(rot_impulse, multiplier=1000))
# print(p5.verts)
# print(get_support(p5.verts, wrap_variable([-1, -1])))
# world = World([p5])
world = World([p1, p2, p3, p4, p5])
run_world(world, screen=self.screen, run_time=180)
def slide_demo(screen):
bodies = []
joints = []
restitution = Defaults.RESTITUTION
fric_coeff = 0.15
inclination = math.pi / 32
r = Rect([inclination, 500, 300], [900, 10],
restitution=restitution, fric_coeff=fric_coeff)
bodies.append(r)
joints.append(TotalConstraint(r))
r = Rect([100, 100], [60, 60],
restitution=restitution, fric_coeff=fric_coeff)
# r = Hull([100, 100], [[30, 30], [-30, 30], [-30, -30], [30, -30]])
bodies.append(r)
r.add_force(Gravity(g=100))
recorder = None
# recorder = Recorder(DT, screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=TIME, screen=screen, recorder=recorder)
def testSlide(self):
bodies = []
joints = []
r = Rect([500, 300], [900, 10])
r.v[0] = math.pi / 32
r.move(1)
r.v[0] = 0.
bodies.append(r)
joints.append(TotalConstraint(r))
r = Rect([100, 100], [60, 60])
r.v[0] = -math.pi / 8 * 0
r.move(1)
r.v[0] = 0.
bodies.append(r)
r.add_force(ExternalForce(down_force, multiplier=100))
# r.add_force(ExternalForce(hor_impulse, multiplier=-100))
# c = Circle([100, 150], 30)
# bodies.append(c)
# c.add_force(ExternalForce(gravity, multiplier=100))
# c = Circle([50, 550], 30)
# c.add_force(ExternalForce(rot_impulse, multiplier=1000))
# bodies.append(c)
# XXX
# c = Circle([875, 100], 30)
# bodies.append(c)
# c.add_force(ExternalForce(gravity, multiplier=100))
recorder = None
# recorder = Recorder(DT, self.screen)
world = World(bodies, joints, dt=DT)
run_world(world, run_time=TIME, screen=self.screen, recorder=recorder)
def testAGrad(self):
def make_world(learned_force):
bodies = []
joints = []
target = Circle([500, 300], 30)
bodies.append(target)
c1 = Circle([250, 210], 30)
bodies.append(c1)
c1.add_force(ExternalForce(learned_force))
c1.add_no_collision(target)
c2 = Circle([400, 250], 30)
bodies.append(c2)
c2.add_no_collision(target)
world = World(bodies, joints, dt=DT)
return world, c2, target
initial_force = torch.DoubleTensor([0, 3, 0])
initial_force[2] = 0
initial_force = Variable(initial_force, requires_grad=True)
# Initial demo
learned_force = lambda t: initial_force if t < 0.1 else ExternalForce.ZEROS
# learned_force = gravity
world, c, target = make_world(learned_force)
# initial_state = world.save_state()
# next_fric_coeff = Variable(torch.DoubleTensor([1e-7]), requires_grad=True)
# c.fric_coeff = next_fric_coeff
# initial_state = world.save_state()
run_world(world, run_time=TIME, screen=None)
learning_rate = 0.001
max_iter = 100
dist_hist = []
last_dist = 1e10
for i in range(max_iter):
learned_force = lambda t: initial_force if t < 0.1 else ExternalForce.ZEROS
world, c, target = make_world(learned_force)
# world.load_state(initial_state)
# world.reset_engine()
# c = world.bodies[0]
# c.fric_coeff = next_fric_coeff
run_world(world, run_time=TIME, screen=None)
dist = (target.pos - c.pos).norm()
dist.backward()
grad = initial_force.grad.data
# grad.clamp_(-10, 10)
initial_force = Variable(initial_force.data - learning_rate * grad,
requires_grad=True)
# grad = c.fric_coeff.grad.data
# grad.clamp_(-10, 10)
# temp = c.fric_coeff.data - learning_rate * grad
# temp.clamp_(1e-7, 1)
learning_rate /= 1.1
# next_fric_coeff = Variable(temp, requires_grad=True)
# print(next_fric_coeff)
if abs((last_dist - dist).data[0]) < 1e-5:
break
last_dist = dist
dist_hist.append(dist)
world = make_world(learned_force)[0]
# c.fric_coeff = next_fric_coeff
# world.load_state(initial_state)
# world.reset_engine()
run_world(world, run_time=TIME, screen=None, recorder=None)
dist = (target.pos - c.pos).norm()
def grad_demo(screen):
initial_force = torch.DoubleTensor([0, 3, 0])
initial_force[2] = 0
initial_force = Variable(initial_force, requires_grad=True)
# Initial demo
learned_force = lambda t: initial_force if t < 0.1 else ExternalForce.ZEROS
# learned_force = gravity
world, c, target = make_world(learned_force)
# initial_state = world.save_state()
# next_fric_coeff = Variable(torch.DoubleTensor([1e-7]), requires_grad=True)
# c.fric_coeff = next_fric_coeff
# initial_state = world.save_state()
run_world(world, run_time=TIME, screen=screen)
learning_rate = 0.001
max_iter = 100
dist_hist = []
last_dist = 1e10
for i in range(max_iter):
learned_force = lambda t: initial_force if t < 0.1 else ExternalForce.ZEROS
world, c, target = make_world(learned_force)
# world.load_state(initial_state)
# world.reset_engine()
# c = world.bodies[0]
# c.fric_coeff = next_fric_coeff
run_world(world, run_time=TIME, screen=None)
dist = (target.pos - c.pos).norm()
dist.backward()
grad = initial_force.grad.data
# grad.clamp_(-10, 10)
initial_force = Variable(initial_force.data - learning_rate * grad,
requires_grad=True)
# grad = c.fric_coeff.grad.data
# grad.clamp_(-10, 10)
# temp = c.fric_coeff.data - learning_rate * grad
# temp.clamp_(1e-7, 1)
learning_rate /= 1.1
# next_fric_coeff = Variable(temp, requires_grad=True)
print(i, '/', max_iter, dist.data[0])
print(grad)
# print(next_fric_coeff)
print(learned_force(0.05))
print('=======')
if abs((last_dist - dist).data[0]) < 1e-5:
break
last_dist = dist
dist_hist.append(dist)
world = make_world(learned_force)[0]
# c.fric_coeff = next_fric_coeff
# world.load_state(initial_state)
# world.reset_engine()
rec = None
# rec = Recorder(DT, screen)
run_world(world, run_time=TIME, screen=screen, recorder=rec)
dist = (target.pos - c.pos).norm()
print(dist.data[0])
# import pickle
# with open('control_balls_dist_hist.pkl', 'w') as f:
# pickle.dump(dist_hist, f)
plot(dist_hist)
def main(screen):
forces = [hor_impulse]
ground_truth_mass = torch.tensor([TOTAL_MASS], dtype=DTYPE)
world, chain = make_world(forces, ground_truth_mass, num_links=NUM_LINKS)
rec = None
# rec = Recorder(DT, screen)
ground_truth_pos = positions_run_world(world, run_time=10, screen=screen, recorder=rec)
ground_truth_pos = [p.data for p in ground_truth_pos]
ground_truth_pos = torch.cat(ground_truth_pos)
learning_rate = 0.5
max_iter = 100
next_mass = torch.rand_like(ground_truth_mass, requires_grad=True)
print('\rInitial mass:', next_mass.item())
print('-----')
optim = torch.optim.RMSprop([next_mass], lr=learning_rate)
loss_hist = []
mass_hist = [next_mass.item()]
last_loss = 1e10
for i in range(max_iter):
if i % 1 == 0:
world, chain = make_world(forces, next_mass.clone().detach(), num_links=NUM_LINKS)
run_world(world, run_time=10, print_time=False, screen=None, recorder=None)
world, chain = make_world(forces, next_mass, num_links=NUM_LINKS)
positions = positions_run_world(world, run_time=10, screen=None)
positions = torch.cat(positions)
positions = positions[:len(ground_truth_pos)]
clipped_ground_truth_pos = ground_truth_pos[:len(positions)]
optim.zero_grad()
loss = MSELoss()(positions, clipped_ground_truth_pos)
loss.backward()
optim.step()
print('Iteration: {} / {}'.format(i+1, max_iter))
print('Loss:', loss.item())
print('Gradient:', next_mass.grad.item())
print('Next mass:', next_mass.item())
print('-----')
if abs((last_loss - loss).item()) < STOP_DIFF:
print('Loss changed by less than {} between iterations, stopping training.'
.format(STOP_DIFF))
break
last_loss = loss
loss_hist.append(loss.item())
mass_hist.append(next_mass.item())
world = make_world(forces, next_mass, num_links=NUM_LINKS)[0]
rec = None
positions = positions_run_world(world, run_time=10, screen=screen, recorder=rec)
positions = torch.cat(positions)
positions = positions[:len(ground_truth_pos)]
clipped_ground_truth_pos = ground_truth_pos[:len(positions)]
loss = MSELoss()(positions, clipped_ground_truth_pos)
print('Final loss:', loss.item())
print('Final mass:', next_mass.item())
plot(loss_hist)
plot(mass_hist)
