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()