def simulate_breakout(a,
world,
ball,
paddle,
blocks,
paddle_idx,
paddle_vel_x=None,
paddle_vel_y=None):
step_breakout(world,
paddle_idx,
a,
paddle_vel_x=paddle_vel_x,
paddle_vel_y=paddle_vel_y)
ball_params = torch.cat([ball.pos, ball.rad.view(1)])
paddle_params = torch.cat([paddle.pos, paddle.dims])
blocks_params = torch.cat([torch.cat([b.pos, b.dims]) for b in blocks])
sn = torch.cat([paddle_params, blocks_params, ball_params])
sn = sn.unsqueeze(0)
new_paddle_params, new_blocks_params, new_ball_params = from_vector(
sn.squeeze(0))
paddle.p = torch.cat([get_tensor([0]), new_paddle_params[0]])
paddle.dims = new_paddle_params[1]
ball.p = torch.cat([get_tensor([0]), new_ball_params[0]])
ball.rad = new_ball_params[1]
for i in range(len(new_blocks_params)):
blocks[i].p = torch.cat([get_tensor([0]), new_blocks_params[i][0]])
blocks[i].dims = new_blocks_params[i][1]
return sn
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 to_vector(paddle_params, blocks_params, ball_params):
ball = torch.cat(
[get_tensor(ball_params[0]),
get_tensor([ball_params[1]])])
paddle = torch.cat([get_tensor(p) for p in paddle_params])
if blocks_params:
blocks = torch.cat([
torch.cat([get_tensor(p) for p in block])
for block in blocks_params
])
else:
# hack: placeholder block outside screen for when there are no blocks
blocks = get_tensor([1000, 60, 1, 1])
ret = torch.cat([paddle, blocks, ball])
if CUDA:
ret = ret.cuda()
return ret
def run_breakout(world,
paddle_idx,
actions,
paddle_vel=None,
dt=0.1,
run_time=10,
screen=None,
recorder=None):
"""Helper function to run a simulation forward once a world is created.
"""
if screen is not None:
import pygame
background = pygame.Surface(screen.get_size())
background = background.convert()
background.fill((255, 255, 255))
elapsed_time = 0.
prev_frame_time = -dt
start_time = time.time()
while world.t < run_time:
a = -0.1
if len(actions) > 0:
a = actions[0]
actions = actions[1:]
v = world.v.clone()
# paddle_vel = 0
# if a == 0:
# paddle_vel = paddle_vel
# elif a == 1:
# paddle_vel = -paddle_vel
v[paddle_idx * 3 + 1] = paddle_vel * a
world.set_v(get_tensor(v))
world.step(fixed_dt=True)
# action = 0
# if len(actions) > 0:
# action = actions[0]
# actions = actions[1:]
# step_breakout(world, paddle_idx, action, paddle_vel)
if screen is not None:
for event in pygame.event.get():
if event.type == pygame.QUIT:
return
if elapsed_time - prev_frame_time >= dt or recorder:
prev_frame_time = elapsed_time
screen.blit(background, (0, 0))
update_list = []
for body in world.bodies:
update_list += body.draw(screen)
for joint in world.joints:
update_list += joint[0].draw(screen)
if not recorder:
# Don't refresh screen if recording
pygame.display.update(update_list)
else:
recorder.record(world.t)
elapsed_time = time.time() - start_time
if not recorder:
# Adjust frame rate dynamically to keep real time
wait_time = dt - (elapsed_time - prev_frame_time)
if wait_time >= 0 and not recorder:
wait_time += dt # XXX
time.sleep(max(wait_time - dt, 0))
elapsed_time = time.time() - start_time
def mpc(s1,
a,
T,
ball_vel=BALL_VEL,
lambda_a=1e-3,
centering=False,
paddle_vel=None,
vert_cost=False,
verbose=1):
ns = len(s1)
na = len(a[0])
class BreakoutDynamics(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, u):
x_dim = x.ndimension()
if x_dim == 1:
x = x.unsqueeze(0)
params = from_vector(x[0])
world, ball, paddle, blocks, paddle_idx = make_world(
*params, ball_vel)
next_x = simulate_breakout(u[0],
world,
ball,
paddle,
blocks,
paddle_idx,
paddle_vel_x=paddle_vel[0],
paddle_vel_y=paddle_vel[1])
return next_x
dynamics = BreakoutDynamics()
if CUDA:
dynamics = dynamics.cuda()
u_lower, u_upper = -ACTION_VAL, ACTION_VAL
x_init = s1.clone()
u_init = get_tensor(a).clone()
if CUDA:
u_init = u_init.cuda()
ball_vel_y = ball_vel[1]
ball_pos_y = s1[-2]
paddle_pos_y = s1[1]
s_Q = torch.zeros(ns)
Q = torch.cat([s_Q, torch.ones(na) * lambda_a
]).type_as(s1).diag().unsqueeze(0).repeat(T, 1, 1)
if vert_cost:
Q[:, ns - 2, ns - 2] = 10
# Simple X tracking
if ball_vel_y > 0:
frames_to_paddle = (paddle_pos_y - ball_pos_y) / ball_vel_y / DT
for t in range(T):
if int(frames_to_paddle) + 1 >= t:
Q[t, ns - 3, ns - 3] = 1
Q[t, ns - 3, 0] = -1
Q[t, 0, 0] = 1
Q[t, 0, ns - 3] = -1
else:
if centering and ball_pos_y > 1000:
Q[:, ns - 3, ns - 3] = 1
Q[:, ns - 3, 0] = -1
Q[:, 0, 0] = 1
Q[:, 0, ns - 3] = -1
p = torch.zeros(ns + na).type_as(Q)
p = p.unsqueeze(0).repeat(T, 1)
Q = Q.unsqueeze(1)
p = p.unsqueeze(1)
x_init = x_init.unsqueeze(0)
solver = MPC(
ns,
na,
T=T,
# x_init=x_init,
u_init=u_init,
u_lower=u_lower,
u_upper=u_upper,
verbose=verbose,
delta_u=10 * ACTION_VAL,
lqr_iter=1,
grad_method=GradMethods.AUTO_DIFF,
n_batch=1,
max_linesearch_iter=1,
exit_unconverged=False,
backprop=False,
)
if CUDA:
solver = solver.cuda()
cost = QuadCost(Q, p)
x, u, objs = solver(x_init, cost, dynamics)
u = u.squeeze(1)
if CUDA:
u = u.cpu()
return u.data.numpy()