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