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)
