class World(PyxelWorld):
def __init__(self):
super().__init__()
self.space = Space(threaded=True)
self.space.collision_slop = 0.7
self.iter = 0
self.time = 0.0
self.on_reset = []
self.P, self.I, self.D = 4.25347222, 0.0001041666, -4.67881944
self.I = 0.0
y = 0
L = 500
self.ground = self.line(-L,
y,
L,
y,
static=True,
friction=0.75,
radius=5)
handler = self.space.add_wildcard_collision_handler(1)
handler.begin = self.on_collision
def update(self):
dt = 1 / 30
with LOCK:
self.space.step(dt)
self.time += dt
def on_collision(self, arbiter, space, data):
for fn in self.on_reset:
fn(self)
return True
def reset(self):
global segway
self.time = 0.0
self.iter += 1
self.clear()
ctrl.object = segway = self.segway()
def clear(self):
for shape in self.space.shapes:
if shape not in self.ground.shapes:
self.space.remove(shape)
for cons in self.space.constraints:
self.space.remove(cons)
for body in self.space.bodies:
if body is not self.ground:
self.space.remove(body)
def segway(self):
R = 12
L = 60
segway = self.circ(0, R, R, mass=1, moment=1, friction=2.0)
arm = self.line(0, R, 0, R + L, mass=2, moment=50)
head = self.circ(0,
R + L,
R,
mass=8,
moment=8,
velocity=(-1, 0),
friction=0.1)
self.pin([head, arm], collide=False)
next(iter(head.shapes)).collision_type = 1
self.pin([segway, arm], anchor=(0, -L / 2), collide=False)
segway.update = iter(self.pid(segway, arm)).__next__
segway.radius = R
segway.arm = arm
segway.head = head
return segway
def pid(self, wheel, arm):
P, I, D = self.P, self.I, self.D
cte = 100
R = wheel.radius
err_ = 0.0
err_i = 0.0
gamma = 0.02
while True:
err = arm.angle
err_i = err + 0.9 * err_i
err_d = err - err_
w = cte * (P * err - I * err_i - D * err_d)
if w == 0:
F = 0
else:
F = 100 * (1 if w > wheel.angular_velocity else -1)
F -= gamma * abs(F) * wheel.angular_velocity
wheel.apply_force_at_local_point((0, +F), (+R, 0))
wheel.apply_force_at_local_point((0, -F), (-R, 0))
err_ = err
yield
class PymunkSimulation:
def __init__(self, do_render, sparse, max_motor_force):
self.do_render = do_render
self.sparse = sparse
self.max_motor_force = max_motor_force
if self.do_render:
pygame.init()
self.screen = pygame.display.set_mode((ENV_SIZE, ENV_SIZE))
self.draw_options = pygame_util.DrawOptions(self.screen)
self.clock = pygame.time.Clock()
self.motors = []
self.segment_bodies = []
self.space = Space()
self.space.iterations = 20
no_collision = self.space.add_collision_handler(NO_COLLISION_TYPE, NO_COLLISION_TYPE)
no_collision.begin = lambda a, b, c: False
ghost_collision = self.space.add_wildcard_collision_handler(GHOST_TYPE)
ghost_collision.begin = lambda a, b, c: False
def set_motor_rates(self, motor_rates):
for i, motor in enumerate(self.motors):
motor.rate = motor_rates[i]
def set_motor_max_forces(self, motor_forces):
for i, motor in enumerate(self.motors):
motor.max_force = motor_forces[i]
def step(self):
dt = 0.03
steps = 30
for i in range(steps):
self.space.step(dt / steps)
def add_arm(self, segment_lengths, anchor_position):
anchor = pymunk.Body(body_type=pymunk.Body.STATIC)
anchor.position = anchor_position
self.space.add(anchor)
segment_anchor = anchor
next_anchor_pos = anchor_position
for i, segment_length in enumerate(segment_lengths):
segment_size = segment_length, 10
segment_body = pymunk.Body(10, pymunk.moment_for_box(10, segment_size))
end_effector_shape = pymunk.Poly.create_box(segment_body, segment_size)
end_effector_shape.collision_type = NO_COLLISION_TYPE
end_effector_shape.friction = 1.0
end_effector_shape.elasticity = 0.1
alpha = random.random() * math.pi * 2
dx = np.cos(alpha) * segment_length / 2
dy = np.sin(alpha) * segment_length / 2
segment_body.position = next_anchor_pos[0] - dx, next_anchor_pos[1] - dy
next_anchor_pos = (next_anchor_pos[0] - 2 * dx, next_anchor_pos[1] - 2 * dy)
segment_body.angle = alpha
anchor_pin_pos = (0 if i == 0 else -segment_lengths[i - 1] / 2, 0)
pin = pymunk.PinJoint(segment_anchor, segment_body, anchor_pin_pos, (segment_length / 2, 0))
self.space.add(pin)
self.space.add(segment_body, end_effector_shape)
motor = pymunk.SimpleMotor(segment_anchor, segment_body, 0)
motor.max_force = self.max_motor_force
self.space.add(motor)
self.motors.append(motor)
self.segment_bodies.append(segment_body)
segment_anchor = segment_body
return segment_anchor, next_anchor_pos
