def step(self, state, action): # print action # returns next state given the action action_map = {0: -0.01, 1: 0, 2: 0.01} state[3] += action_map[action.index(1)] state_dictionary = {} for state_name, state_element in zip(self.state_dim_names, state): state_dictionary[state_name] = state_element next_state_dictionary = simulator2d.simulate_timestep(state_dictionary, self.dt) next_state = [] for name in self.state_dim_names: next_state.append(next_state_dictionary[name]) self.state = next_state return next_state
def display_configuration(state): screen.fill(black) servo_end = rotating_mass_endpoint(state) arm_end = arm_endpoint(state) points = [toscreen((0, 0)), toscreen(arm_end)] pygame.draw.lines(screen, red, False, points, 2) points = [toscreen(arm_end), toscreen(servo_end)] pygame.draw.lines(screen, green, False, points, 2) mg_pt = (state["cgx"], -state["cgy"]) points = [toscreen(scale_pt(mg_pt)), toscreen(scale_pt(mg_pt, 3))] pygame.draw.lines(screen, white, False, points, 10) pygame.display.update() direction = 1 while True: for event in pygame.event.get(): if event.type == pygame.QUIT: pygame.quit() sys.exit() display_configuration(s) pygame.time.delay(10) s = sim.simulate_timestep(s, 0.1)