def get_obs(self):
obs = []
for state in self.extra_data.states:
body = find_body(self.world, state.body)
if state.typ == "xpos":
obs.append(body.position[0])
elif state.typ == "ypos":
obs.append(body.position[1])
elif state.typ == "xvel":
obs.append(body.linearVelocity[0])
elif state.typ == "yvel":
obs.append(body.linearVelocity[1])
elif state.typ == "apos":
obs.append(body.angle)
elif state.typ == "avel":
obs.append(body.angularVelocity)
else:
raise NotImplementedError
return np.array(obs)
def step(self, state, action):
if len(action) != self.action_dim:
raise ValueError('incorrect action dimension: expected %d but got '
'%d' % (self.action_dim, len(action)))
if not np.array_equal(state, self.current_state):
self.set_state(state)
for ctrl, act in zip(self.extra_data.controls, action):
if ctrl.typ == "force":
body = find_body(self.world, ctrl.body)
direction = np.array(ctrl.direction)
direction = direction / np.linalg.norm(direction)
world_force = body.GetWorldVector(direction * act)
world_point = body.GetWorldPoint(ctrl.anchor)
body.ApplyForce(world_force, world_point, wake=True)
else:
raise NotImplementedError
self.world.Step(
self.extra_data.timeStep,
self.extra_data.velocityIterations,
self.extra_data.positionIterations
)
return self.get_state(), self.get_obs()
