from collections import OrderedDict import numpy as np from visual_dynamics.envs import RosEnv try: import rospy from visual_dynamics.pr2 import camera_sensor from visual_dynamics.pr2 import PR2 except: pass from visual_dynamics.spaces import BoxSpace, TupleSpace class Pr2Env(RosEnv): def __init__(self, action_space, observation_space, state_space, sensor_names, dt=None): super(Pr2Env, self).__init__(action_space, observation_space, state_space, sensor_names) self._dt = 0.2 if dt is None else dt self.pr2 = PR2.PR2() self.pr2.larm.goto_posture('side') self.pr2.rarm.goto_posture('side') self.pr2.torso.go_down() gains = {'head_pan_joint': {'d': 2.0, 'i': 12.0, 'i_clamp': 0.5, 'p': 50.0}, 'head_tilt_joint': {'d': 3.0, 'i': 4.0, 'i_clamp': 0.2, 'p': 1000.0}} rospy.set_param('/head_traj_controller/gains', gains) self.pr2.head.set_pan_tilt(*((self.state_space.low + self.state_space.high) / 2.0)) self.msg_and_camera_sensor = camera_sensor.MessageAndCameraSensor() rospy.sleep(5.0) def step(self, action): # update action to be within the action space if not self.action_space.contains(action): action = self.action_space.clip(action, out=action) pan_tilt_angles = self.pr2.head.get_joint_positions() action[:] = self.state_space.clip(pan_tilt_angles + action) - pan_tilt_angles self.pr2.head.command_pan_tilt_vel(*action, dt=self._dt) rospy.sleep(.3) def get_state(self): return self.pr2.head.get_joint_positions() def reset(self, state=None): if state is None: state = self.state_space.sample() self.pr2.head.goto_joint_positions(state) rospy.sleep(1.0) def get_error_names(self): return ['pan_angle', 'tilt_angle'] def get_errors(self, target_state): pan_error, tilt_error = np.abs(target_state - self.get_state()) return OrderedDict([('pan_angle', pan_error), ('tilt_angle', tilt_error)]) def observe(self): _, obs = self.get_state_and_observe() return obs def get_state_and_observe(self): joint_state_msg, image = self.msg_and_camera_sensor.get_msg_and_observe() state = self.pr2.head.get_joint_positions(msg=joint_state_msg) obs = [] for sensor_name in self.sensor_names: if sensor_name == 'image': observation = image else: raise ValueError('Unknown sensor name %s' % sensor_name) obs.append(observation.copy()) return state, obs def render(self): pass @property def dt(self): return self._dt def _get_config(self): config = super(Pr2Env, self)._get_config() return config def main(): rospy.init_node('camera_sensor', anonymous=True) action_space = BoxSpace(np.deg2rad([-5., -5.]), np.deg2rad([5., 5.])) observation_space = TupleSpace([spaces.BoxSpace(0, 255, shape=(240, 320, 3), dtype=np.uint8)]) state_space = BoxSpace(np.deg2rad([-30., 45.]), np.deg2rad([30., 75.])) sensor_names = ['image'] pr2_env = Pr2Env(action_space, observation_space, state_space, sensor_names) import IPython as ipy; ipy.embed() if __name__ == "__main__": main()