def _bg_loop_fn(self):
if not self._running:
# publish steady hold message when asking for zero motion
motion = CFMotion()
motion.is_flow_motion = True
motion.stamp.stamp = rospy.Time.now()
self._ros_motion_pub.publish(motion)
def get_action(self, step, current_episode_step, observation, goal,
explore):
# chosen_actions, chosen_values, action_infos = self.get_actions([step], [current_episode_step] [observation],
# [goal], explore=explore)
# return chosen_actions[0], chosen_values[0], action_infos[0]
motion = CFMotion()
motion.yaw = 0
# return motion
print("TELEOP POLICY got action:",
[motion.x, motion.y, motion.yaw, motion.dz])
return [motion.yaw], None, dict()
def get_action(self, step, current_episode_step, observation, goal,
explore):
# chosen_actions, chosen_values, action_infos = self.get_actions([step], [current_episode_step] [observation],
# [goal], explore=explore)
# return chosen_actions[0], chosen_values[0], action_infos[0]
motion = CFMotion()
motion.yaw = random.uniform(-120, 120)
# return motion
print("RANDOM POLICY chose action:",
[motion.x, motion.y, motion.yaw, motion.dz])
return [motion.yaw], None, dict()
def _reset():
global _is_collision
_is_collision = False
global _joy_stop_trash_btn_pressed
global _joy_stop_save_btn_pressed
global _joy_estop_trash_btn_pressed
global _joy_estop_pause_btn_pressed
global _curr_motion
_joy_stop_trash_btn_pressed = False
_joy_stop_save_btn_pressed = False
_joy_estop_trash_btn_pressed = False
_joy_estop_pause_btn_pressed = False
_curr_motion = CFMotion()
_curr_motion.is_flow_motion = True
def _set_motion(self, vx, vy, vz):
motion = CFMotion()
motion.x = vx
motion.y = vy
motion.yaw = 0 # yaw disabled for now
motion.dz = vz # convert velocity to a dt for this time step
motion.is_flow_motion = True
self._curr_motion = motion
def _set_motion(self, vx, vy, vz):
motion = CFMotion()
motion.x = vx
motion.y = vy
motion.yaw = 0 # yaw disabled for now
motion.dz = vz
motion.is_flow_motion = True
self._curr_motion = motion
def _background_thread():
global _curr_motion
rate = rospy.Rate(10)
while not rospy.is_shutdown():
if _rosbag.is_open:
_curr_motion.stamp.stamp = rospy.Time.now()
_motion_pub.publish(_curr_motion)
# drop things into the rosbag
# _queue_lock.acquire()
# _rosbag.write_all(_ros_msg_queue)
# _buffer_lock.acquire()
# _swap()
# _queue_lock.release()
# _buffer_lock.release()
else:
temp_motion = CFMotion()
temp_motion.stamp.stamp = rospy.Time.now()
if _curr_joy.axes:
temp_motion.y = _curr_joy.axes[ROLL_AXIS] * VY_SCALE
temp_motion.x = _curr_joy.axes[PITCH_AXIS] * VX_SCALE
temp_motion.yaw = _curr_joy.axes[YAW_AXIS] * YAW_SCALE
temp_motion.is_flow_motion = True
_motion_pub.publish(temp_motion)
rate.sleep()
def compute_motion(self):
#pulls latest joystick data
if not self.use_joy:
# no motion input from controller
return None
motion = None
self.access_lock.acquire()
if self.cmd != -1:
motion = CFCommand()
if self.cmd == CFCommand.ESTOP:
motion.cmd = CFCommand.ESTOP
elif self.cmd == CFCommand.TAKEOFF:
motion.cmd = CFCommand.TAKEOFF
elif self.cmd == CFCommand.LAND:
motion.cmd = CFCommand.LAND
#reset
self.cmd = -1
#repeat send at 10Hz
elif self.curr_joy:
motion = CFMotion()
motion.is_flow_motion = self.is_flow_motion
# computing regular vx, vy, yaw, alt motion
if self.is_flow_motion:
motion.y = self.curr_joy.axes[ROLL_AXIS] * VY_SCALE
motion.x = self.curr_joy.axes[PITCH_AXIS] * VX_SCALE
else:
motion.x = self.curr_joy.axes[ROLL_AXIS] * ROLL_SCALE
motion.y = self.curr_joy.axes[PITCH_AXIS] * PITCH_SCALE
motion.x = motion.x if abs(motion.x) > ROLL_CLIP else 0
motion.y = motion.y if abs(motion.y) > PITCH_CLIP else 0
#common
motion.yaw = self.curr_joy.axes[YAW_AXIS] * YAW_SCALE
# print(self.curr_joy.axes)
motion.dz = self.curr_joy.axes[THROTTLE_AXIS] * THROTTLE_SCALE
# print("ALT CHANGE: %.3f" % motion.dz)
self.access_lock.release()
return motion
def reset(self, ret_latent=False):
self._curr_motion = CFMotion()
self._curr_motion.is_flow_motion = True
vec, latent_idx = self._sim.reset(
[random.random(),
random.random(), 0.5 * random.random()],
ret_latent=True)
self._latent_idx = latent_idx
self._obs[0] = vec.vector.x
self._obs[1] = vec.vector.y
self._obs[2] = vec.vector.z
self.render_clear()
obs, goal = self.reset_model()
if ret_latent:
return obs, goal, latent_idx
else:
return obs, goal
class Controller:
DO_NOTHING_CMD = CFMotion()
COLLISION_THRESH = 0.5
def __init__(self, ID):
self.id = ID
self.mat = None
self.data = None
#need to facilitate a set of publishers per cf node
self.data_sub = rospy.Subscriber('cf/%d/data' % ID, CFData,
self.data_cb)
self.image_sub = rospy.Subscriber('cf/%d/image' % ID, CompressedImage,
self.image_cb)
self.cmd_pub = rospy.Publisher('cf/%d/command' % self.id,
CFCommand,
queue_size=10)
self.motion_pub = rospy.Publisher('cf/%d/motion' % self.id,
CFMotion,
queue_size=10)
self.coll_pub = rospy.Publisher('cf/%d/coll' % self.id,
Bool,
queue_size=10)
def compute_motion(self):
print("Doing nothing -- ")
return None
def observedCollision(self):
# if not self.data:
# return False
# xy_mag = math.sqrt(self.data.accel_x ** 2 + self.data.accel_y ** 2)
# print("-- COLLIDED --") if xy_mag > self.COLLISION_THRESH else None
# return xy_mag > self.COLLISION_THRESH
return False
## CALLBACKS ##
def image_cb(self, msg):
pil_jpg = io.BytesIO(msg.data)
pil_arr = Image.open(pil_jpg).convert('L') #grayscale
self.mat = np.array(pil_arr)
pass
def data_cb(self, msg):
self.data = msg
pass
## SETTERS ##
## THREADS ##
def run(self):
rate = rospy.Rate(10) # 10Hz
while not rospy.is_shutdown():
action = self.compute_motion()
if self.data and self.observedCollision():
print("-- COLLISION : E-stopping --")
action = CFCommand()
action.cmd = CFCommand.ESTOP
action.stamp.stamp = rospy.Time.now()
self.cmd_pub.publish(action)
self.coll_pub.publish(True)
#sleep for 2 seconds
rospy.Rate(1.0 / 2).sleep()
print("Back Online")
else:
self.coll_pub.publish(False)
if isinstance(action, CFMotion):
# if action != Controller.DO_NOTHING_CMD:
action.stamp.stamp = rospy.Time.now()
self.motion_pub.publish(action)
# else:
# print("--- DO NOTHING CMD SENT ---")
elif isinstance(action, CFCommand):
action.stamp.stamp = rospy.Time.now()
self.cmd_pub.publish(action)
print("CALLED COMMAND -> %s" % cmd_type[action.cmd])
else:
pass
# rospy.spinOnce()
rate.sleep()
def __init__(self, params, env_spec):
env_spec.assert_has_names(['obs', 'act', 'prev_obs', 'prev_act'])
self._env_spec = env_spec
self.x_dim = self._env_spec.names_to_shapes.obs[-1] # dimension of obs
self.u_dim = self._env_spec.names_to_shapes.act[-1] # dimension of act
self._obs_hist_len = self._env_spec.names_to_shapes.prev_obs[-2]
self._act_hist_len = self._env_spec.names_to_shapes.prev_act[-2]
# Setup the parameters
self._dt = params.dt
self._ros_prefix = params.ros_prefix
self.offline = params.offline # running offline means don't publish anything (e.g. from rosbag)
self.horizon = int(params.horizon)
self._normalize = params.normalize
if self._normalize:
logger.debug("Normalizing input ON")
self._coef_cx = 1. / params.img_width
self._coef_cy = 1. / params.img_height
self._coef_area = 1. / (params.img_width * params.img_height)
else:
self._coef_cx = self._coef_cy = self._coef_area = 1.
self.time_last_step = -1
self._obs = np.zeros((self.x_dim,))
self._next_obs = np.zeros((self.x_dim,))
self._prev_obs = None
self._prev_act = None
self._ob_lb, self._ob_ub = self._env_spec.limits(['obs'])
self._ob_lb, self._ob_ub = self._ob_lb[0], self._ob_ub[0]
self.observation_space = spaces.Box(self._ob_lb, self._ob_ub)
self._initial_goal_pos = params.initial_goal_pos
self._curr_goal_pos = None
self._cost = 0
self._done = False # indicates if the copter has completed/failed its rollout
self._running = False # indicates if a rollout is currently running
# start with an assumption of collision to allow for takeoff
self.collided = True
self._prev_act = None
self._ac_lb, self._ac_ub = self._env_spec.limits(['act'])
self._ac_lb, self._ac_ub = self._ac_lb[0], self._ac_ub[0]
self.action_space = spaces.Box(self._ac_lb, self._ac_ub)
# ros callbacks
self._curr_motion = CFMotion()
self._ros_topics_and_types = {
self._ros_prefix + 'coll': Bool,
self._ros_prefix + 'data': CFData,
'extcam/target_vector': Vector3Stamped,
}
self._ros_msgs = dict()
self._ros_msg_times = dict()
# ROS setup
ru.setup_ros_node('MBRLNode', anonymous=True)
for topic, type in self._ros_topics_and_types.items():
ru.bind_subscriber(RosTopic(topic, type), self.ros_msg_update, callback_args=(topic,))
# does not use these publishers if _offline
_, self._ros_motion_pub = ru.get_publisher(RosTopic(self._ros_prefix + "motion", CFMotion), queue_size=10)
_, self._ros_command_pub = ru.get_publisher(RosTopic(self._ros_prefix + "command", CFCommand), queue_size=10)
_, self._ros_goal_pub = ru.get_publisher(RosTopic("/mpc/goal_vector", Vector3Stamped), queue_size=10)
_, self._ros_action_pub = ru.get_publisher(RosTopic("/mpc/action_vector", Vector3Stamped), queue_size=10)
_, self._ros_reward_pub = ru.get_publisher(RosTopic("/mpc/reward_vector", Vector3Stamped), queue_size=10)
_, self._ros_latent_pub = ru.get_publisher(RosTopic("/mpc/latent_vector", Vector3Stamped), queue_size=10)
_, self._ros_trajectory_marker_pub = ru.get_publisher(RosTopic("/mpc/trajectory_marker", Marker), queue_size=500)
_, self._ros_ac_marker_pub = ru.get_publisher(RosTopic("/mpc/action_marker", Marker), queue_size=10)
logger.info("[COPTER]: Initialized, waiting for topic streams before starting...")
while not self.ros_is_good(display=False):
pass
logger.info("[COPTER]: Topics are active.")
self._step_rate = rospy.Rate(1. / self._dt) # used in step extra function
self._rate = rospy.Rate(1. / self._dt) # used only in background thread
# this background publisher thread is only when we are online
if not self.offline:
ru.setup_background_thread(self._bg_loop_fn, self._rate)
ru.start_background_thread()
self.reset()
def _step(self, u, no_rew_pub=False, step_extra_func=None, buffer_window_size=-1, **kwargs):
self._running = True # signifies that we are publishing actions now
if not self.offline:
assert self.ros_is_good(display=True)
# else:
# u = self._latest_offline_ac.copy()
obs = self._next_obs.copy()
self._enforce_dimension(obs, u)
# Limit the control inputs
u0 = np.clip(u, self._ac_lb, self._ac_ub)
# update vx, vy, dz (does NOT publish)
self._set_motion(u0[0], u0[1], u0[2])
self._done = np.all(obs == self._curr_goal_pos) or self.collided
if np.any(np.isnan(obs)):
logger.debug('[CF]: NAN POSITION')
# publishing action messages happens only if we are fully online
if not self.offline:
if any((self._curr_motion.x, self._curr_motion.y, self._curr_motion.yaw,
self._curr_motion.dz)) or not self.collided:
self._curr_motion.stamp.stamp = rospy.Time.now()
self._ros_motion_pub.publish(self._curr_motion)
else:
# publish steady hold message when asking for zero motion
motion = CFMotion()
motion.is_flow_motion = True
motion.stamp.stamp = rospy.Time.now()
self._ros_motion_pub.publish(motion)
else:
# offline action publishing
act_vec = Vector3Stamped()
act_vec.header.stamp = rospy.Time.now()
act_vec.vector.x = self._curr_motion.x
act_vec.vector.y = self._curr_motion.y
act_vec.vector.z = self._curr_motion.dz
self._ros_action_pub.publish(act_vec)
goal_vec = Vector3Stamped()
goal_vec.header.stamp = rospy.Time.now()
goal_vec.vector.x = self._curr_goal_pos[0]
goal_vec.vector.y = self._curr_goal_pos[1]
goal_vec.vector.z = self._curr_goal_pos[2]
self._ros_goal_pub.publish(goal_vec)
# TODO: set cost properly and in a single place
if not no_rew_pub:
self.pub_rew(-self._cost)
## EXTRA FUNCTIONALITY for inferring latent online
ret = None
if step_extra_func is not None:
def exit_condition(iters): return self._step_rate.remaining().to_sec() < 1e-7 \
and self._step_rate.sleep() is None
ret = step_extra_func(exit_cond=exit_condition,
buffer_window_size=buffer_window_size) # returns at or before sleep_end_time
mu = ret[0][0].item()
sig = ret[1][0].exp().item() # TODO do something with these
latent_vec = Vector3Stamped()
latent_vec.header.stamp = rospy.Time.now()
latent_vec.vector.x = mu
latent_vec.vector.y = sig
self._ros_latent_pub.publish(latent_vec)
else:
self._step_rate.sleep() # sleep until next iter
# history update
self._prev_obs = advance_history_np(self._prev_obs[None], obs[None])[0]
self._prev_act = advance_history_np(self._prev_act[None], u0[None])[0]
self._next_obs = self._obs.copy() # the exact obs the policy sees
return self.get_obs(), self.get_goal(), self._done
def joy_cb(msg):
global _prev_joy
global _curr_joy
global _joy_stop_trash_btn_pressed
global _joy_stop_save_btn_pressed
global _joy_estop_trash_btn_pressed
global _joy_estop_pause_btn_pressed
global _curr_motion
_prev_joy = _curr_joy
_curr_joy = msg
# print(_curr_joy.buttons)
# print(_prev_joy.buttons)
# print()
if _is_btn(_joy_stop_trash_btn):
_joy_stop_trash_btn_pressed = True
if _is_btn(_joy_stop_save_btn):
_joy_stop_save_btn_pressed = True
if _is_btn(_joy_estop_trash_btn):
_joy_estop_trash_btn_pressed = True
if _is_btn(_joy_estop_pause_btn):
_joy_estop_pause_btn_pressed = True
# need to translate curr motion
_curr_motion = CFMotion()
_curr_motion.is_flow_motion = _is_flow_motion
# computing regular vx, vy, yaw, alt motion
if _is_flow_motion:
_curr_motion.y = _curr_joy.axes[ROLL_AXIS] * VY_SCALE
_curr_motion.x = _curr_joy.axes[PITCH_AXIS] * VX_SCALE
else:
_curr_motion.x = _curr_joy.axes[ROLL_AXIS] * ROLL_SCALE
_curr_motion.y = _curr_joy.axes[PITCH_AXIS] * PITCH_SCALE
_curr_motion.x = _curr_motion.x if abs(_curr_motion.x) > ROLL_CLIP else 0
_curr_motion.y = _curr_motion.y if abs(_curr_motion.y) > PITCH_CLIP else 0
#common
_curr_motion.yaw = _curr_joy.axes[YAW_AXIS] * YAW_SCALE
_curr_motion.dz = 0
# parser.add_argument('-logdir', type=str, required=True)
args = parser.parse_args()
# main(args.rosbag_dir, args.ros_prefix, args.is_flow, args.max_steps)
rospy.init_node("DataCapture", anonymous=True)
_rosbag_lock = threading.Lock()
_queue_lock = threading.Lock()
# _buffer_lock = threading.Lock()
_ros_prefix = args.ros_prefix
_is_flow_motion = args.is_flow
_rosbag = RolloutRosbag(args.rosbag_dir)
_is_collision = False
_curr_motion = CFMotion()
_curr_joy = Joy()
_prev_joy = Joy()
_joy_start_btn = 1 # A
_joy_stop_trash_btn = 3 # Y
_joy_stop_save_btn = 0 # X
_joy_estop_trash_btn = 2 # B
_joy_estop_pause_btn = 5 # RB
_joy_stop_trash_btn_pressed = False
_joy_stop_save_btn_pressed = False
_joy_estop_trash_btn_pressed = False
_joy_estop_pause_btn_pressed = False
_ros_msg_times = dict()