mpc_joint_node.cpp

Go to the documentation of this file.

#include <ros/ros.h>

#include <boost/thread/mutex.hpp>

#include <tf/tf.h>
#include <tf/transform_listener.h>
#include <tf/tfMessage.h>
#include <tf/transform_broadcaster.h>

#include <sensor_msgs/JointState.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/QuaternionStamped.h>
#include <trajectory_msgs/JointTrajectory.h>
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>

#include <pr2_controllers_msgs/Pr2GripperCommandAction.h>
#include <actionlib/client/simple_action_client.h>

#include <motld/TrackedObjects.h>

#include <mpc/types.h>
#include <mpc/mpc.h>
#include <mpc/finite_differences.h>
#include <reactive_grasping_pr2/pr2_joint_space.h>

#ifndef SQR
#define SQR(X) ((X)*(X))
#endif

using namespace MPC;

typedef actionlib::SimpleActionClient<pr2_controllers_msgs::Pr2GripperCommandAction> GripperClient;

class ArmMPCJoints {

private:
        bool initialized_;
        VectorT center_;
        tf::TransformBroadcaster *tf_broadcaster_;
        tf::TransformListener  *tf_transformer_;

        ros::Subscriber obj_sub_;
        
        ros::Publisher joint_pose_pub_;
        ros::Publisher marker_pub_;
        ros::Subscriber joint_state_sub_;
        boost::mutex state_mutex_;
        std::vector<std::string> joint_names_;
        std::vector<int> joint_mapping_;
        bool joint_mapping_initialized_;
        VectorT curr_joint_state_;

        int curr_t_;
        float eps_;
        float max_mov_dist_;

        bool on_way_to_goal_;
        bool last_goal_set_;
        bool last_pos_set_;
        VectorT last_goal_;
        VectorT last_pos_;

        int n, m, N;
        bool restart_ddp_;
        DDPParams params;
        Trajectory curr_traj_;
        MatrixT u0_;
        trajectory_msgs::JointTrajectory joint_traj_;
        GripperClient *close_gripper_;

        pr2_joint_space::RightArmModel *r_arm_model_;

        ros::NodeHandle nh_;
        ros::NodeHandle private_nh_;

        // names of parameters to read
        std::string root_frame_;
        std::string gripper_frame_;
        std::string object_topic_;
        std::string gripper_action_topic_;
        std::string gripper_close_topic_;
        std::string marker_topic_;

public:
        ArmMPCJoints(ros::NodeHandle &nh) : nh_(nh), private_nh_(ros::NodeHandle("~")), initialized_(false) {
                // read parameters
                private_nh_.param("root_frame", root_frame_, std::string("/torso_lift_link"));
                private_nh_.param("gripper_frame", gripper_frame_, std::string("/r_wrist_roll_link"));
                private_nh_.param("object_topic", object_topic_, std::string("/objects"));
                private_nh_.param("gripper_action_topic", gripper_action_topic_, std::string("/r_arm_controller/command"));
                private_nh_.param("gripper_close_topic", gripper_close_topic_, std::string("/r_gripper_controller/gripper_action"));
                private_nh_.param("marker_topic", marker_topic_, std::string("MPC_markers"));

                // setup tf transformer and broadcaster
                tf_broadcaster_ = new tf::TransformBroadcaster();
                tf_transformer_ = new tf::TransformListener();

                // setup mpc callback
                obj_sub_ = nh_.subscribe("objects", 1, &ArmMPCJoints::mpc_callback, this);
                // setup pose client
                joint_pose_pub_ = nh_.advertise<trajectory_msgs::JointTrajectory>(gripper_action_topic_, 1);
                // setup gripper client
                close_gripper_ = new GripperClient(gripper_close_topic_, true);
                while(!close_gripper_->waitForServer(ros::Duration(5.0))) {
                        ROS_INFO("Waiting for the gripper action server to come up");
                }                     

                // setup visualization
                marker_pub_ = nh_.advertise<visualization_msgs::MarkerArray>(marker_topic_, 1); 

                // setup DDP and dynamics
                n = 14;
                m = 7;
                N = 20;
                u0_.setZero(m, N);
                VectorT internal_goal(VectorT::Zero(14));
                // setup some constants
                eps_ = 0.02;
                max_mov_dist_ = 0.05;
                // plan full number of steps in the first callback
                restart_ddp_ = true;

                // real euclidean robot dynamics
                r_arm_model_ = new pr2_joint_space::RightArmModel(nh_, root_frame_);
                r_arm_model_->setInternalGoal(internal_goal);
                
                finite_differences::setBaseFunctions(std::bind(&pr2_joint_space::RightArmModel::dynamics, r_arm_model_, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3), std::bind(&pr2_joint_space::RightArmModel::cost, r_arm_model_, std::placeholders::_1, std::placeholders::_2));
                bool init_success = DDP_init(std::bind(&pr2_joint_space::RightArmModel::dynamics, r_arm_model_, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3), std::bind(&pr2_joint_space::RightArmModel::dynamicsD, r_arm_model_, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3), &finite_differences::cost, &finite_differences::costD);
                //bool init_success = DDP_init(std::bind(&pr2_joint_space::RightArmModel::dynamics, r_arm_model_, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3), &finite_differences::dynamicsD, &finite_differences::cost, &finite_differences::costD);
                //*/
                // setup one obstacle
                /*
                VectorT otmp(VectorT::Zero(n));
                obs_ = &pr2_arm_ik::getObstacles();
                obs_->clear();
                obs_->push_back(otmp);
                */
                // init temporaries
                center_.resize(3);
                last_goal_.resize(n);
                last_pos_.resize(n);
                on_way_to_goal_ = false;
                last_goal_set_ = false;
                last_pos_set_ = false;
                curr_t_ = 1;
                // name the joints that we want to control
                // create a mapping for the joints we want to control
                joint_names_.resize(7);
                joint_names_[0] = "r_shoulder_pan_joint";
                joint_names_[1] = "r_shoulder_lift_joint";
                joint_names_[2] = "r_upper_arm_roll_joint";
                joint_names_[3] = "r_elbow_flex_joint";
                joint_names_[4] = "r_forearm_roll_joint";
                joint_names_[5] = "r_wrist_flex_joint";
                joint_names_[6] = "r_wrist_roll_joint";

                // create a mapping for the joints 
                // it will be initialized the first time a joint_state
                // message is received
                joint_mapping_.resize(7);
                joint_mapping_initialized_ = false;
                // prepare joint trajectory message
                joint_traj_.joint_names.resize(joint_names_.size());
                for (int i = 0; i < joint_names_.size(); ++i) {
                        joint_traj_.joint_names[i] = joint_names_[i];
                }
                joint_traj_.points.resize(1);

                joint_traj_.points[0].positions.resize(7);
                joint_traj_.points[0].velocities.resize(7);
                for (size_t j = 0; j < 7; ++j) {
                        joint_traj_.points[0].velocities[j] = 0.0;
                }
                // init the joint_state subscriber
                joint_state_sub_ = nh_.subscribe("/joint_states", 1, &ArmMPCJoints::joint_state_callback, this);
                curr_joint_state_.setZero(joint_names_.size()*2);

                // mark as initialized
                initialized_ = init_success;
                ROS_INFO("Done initializing");
        }

        ~ArmMPCJoints() {
                delete obj_sub_;
                delete close_gripper_;
        }

        void open_gripper(bool wait=false) {
                pr2_controllers_msgs::Pr2GripperCommandGoal open;
                open.command.position = 0.08;
                open.command.max_effort = -1.0;
    
                close_gripper_->sendGoal(open);
                if (wait) {
                        close_gripper_->waitForResult();
                        if(close_gripper_->getState() == actionlib::SimpleClientGoalState::SUCCEEDED)
                                ROS_INFO("The gripper opened!");
                        else
                                ROS_INFO("The gripper failed to open.");
                }
        } 
        
        void close_gripper(bool wait=false) {
                pr2_controllers_msgs::Pr2GripperCommandGoal close;
                close.command.position = 0.0;
                close.command.max_effort = 20.0;
    
                close_gripper_->sendGoal(close);
                if (wait) {
                        close_gripper_->waitForResult();
                        if(close_gripper_->getState() == actionlib::SimpleClientGoalState::SUCCEEDED)
                                ROS_INFO("The gripper closed!");
                        else
                                ROS_INFO("The gripper failed to close.");
                }
        } 

        void joint_state_callback(const sensor_msgs::JointStateConstPtr &msg){
                boost::mutex::scoped_lock(state_mutex_);
                if (!joint_mapping_initialized_) {
                        for (int  i = 0; i < joint_names_.size(); ++i) {
                                for (int j = 0; j < msg->name.size(); ++j) {
                                        if (joint_names_[i].compare(msg->name[j]) == 0) { 
                                                joint_mapping_[i] = j;
                                        }
                                }
                        }
                        joint_mapping_initialized_ = true;
                }
                
                for (int  i = 0; i < joint_mapping_.size(); ++i) {
                        curr_joint_state_(i) = msg->position[joint_mapping_[i]];
                        curr_joint_state_(i+joint_mapping_.size()) = msg->velocity[joint_mapping_[i]];
                }
        }

        void get_cart_and_joint_pos(VectorT &x0) {
                if (!joint_mapping_initialized_) {
                        x0.setZero();
                        return;
                }
                {
                        boost::mutex::scoped_lock(state_mutex_);
                        // get last saved joint state
                        x0.tail(curr_joint_state_.size()) = curr_joint_state_;
                }
        }


        void mpc_callback(const motld::TrackedObjectsConstPtr &objects) {
                ROS_INFO("mpc_callback called!");
                if (not initialized_) {
                        ROS_WARN("mpc not initialized");
                        return;
                }
                if (objects->name.size() < 1) {
                        ROS_WARN("no object poses received");
                        return;
                }
                if (objects->name.size() != objects->pose.poses.size()) {
                        ROS_WARN("object poses and object names array sizes do not match");
                        return;
                }
                
                if (on_way_to_goal_ && curr_t_ +1 >= curr_traj_.x.cols()) {
                        // TODO: there was a sleep here in the python implementation to prevent weird near goal behaviour 
                        // when the last action has not been completely executed
                }

                ros::Time goal_stamp = objects->pose.header.stamp;
                ros::Time time_stamp_now = ros::Time::now();
                // we always assume that the first object is the goal
                VectorT goal(3);
                goal(0) = objects->pose.poses[0].position.x;
                goal(1) = objects->pose.poses[0].position.y;
                goal(2) = objects->pose.poses[0].position.z;

                if (tf_transformer_->waitForTransform(objects->pose.header.frame_id, gripper_frame_, goal_stamp, ros::Duration(1.)) == false) {
                        ROS_ERROR("Did not find transform between %s and %s", objects->pose.header.frame_id.c_str(), gripper_frame_.c_str());
                        return;
                }
                ROS_INFO("Got transform!");

                tf::StampedTransform trans;
                tf_transformer_->lookupTransform(objects->pose.header.frame_id, gripper_frame_, goal_stamp, trans);

                // set new center
                center_(0) = trans.getOrigin().x();
                center_(1) = trans.getOrigin().y();
                center_(2) = trans.getOrigin().z();
                
                // set external (euclidean) goal pos
                r_arm_model_->setExternalGoal(goal);

                VectorT x0(14);
                get_cart_and_joint_pos(x0);
                VectorT curr_joint_goal = x0;
                r_arm_model_->ik(center_, curr_joint_goal);
                curr_joint_goal.tail(7).setZero(); // assume 0 vel at goal
                r_arm_model_->setInternalGoal(curr_joint_goal);
                // internal coordinates
                center_ -= goal;

                /*
                (*obs_)[0](0) = obstacle->point.x - goal->point.x;
                (*obs_)[0](1) = obstacle->point.y - goal->point.y;
                (*obs_)[0](2) = obstacle->point.z - goal->point.z;
                */
                float goal_diff = center_.norm();

                if (goal_diff < eps_ + 0.02) {
                        ROS_INFO("already at goal / goal reached");
                        close_gripper();
                        return;
                }
                // TODO: insert obstacles
                float goal_change = 200.;
                float pos_change = 100.;
                if (last_goal_set_) {
                        goal_change = (last_goal_ - goal).norm();
                }
                if (last_pos_set_) {
                        pos_change = (center_ - last_pos_).norm();
                }
                ROS_INFO("pchange: %f gchange: %f max_dist: %f", pos_change, goal_change, max_mov_dist_);
                if (pos_change > max_mov_dist_ || goal_change > max_mov_dist_) {
                        // reset policy, since we are too far from our last computed one
                        //actions_.setZero(actions_.rows(), actions_.cols());
                        on_way_to_goal_ = false;
                        restart_ddp_ = true;
                }

                /* check if we have a policy that gets us to the goal
                   if so and the goal did not change significantly execute it */
                int x_len = curr_traj_.x.cols();
                // TODO: check if we want to enable the on_way_to_goal_ behavior or not!
                if (false && on_way_to_goal_ && goal_change < eps_ && curr_t_ < x_len - 1) {
                        curr_t_++;
                } else {
                        // run ddp
                        if (restart_ddp_) {
                                ROS_INFO("restarting ddp");
                                u0_.setZero();
                                DDP(x0, u0_, params, 100, curr_traj_);
                                restart_ddp_ = false;
                        } else {
                                u0_ = curr_traj_.u;
                                DDP(x0, u0_, params, 10, curr_traj_);
                        }
                        curr_t_ = 1;
                }

                // fill gripper message
                const VectorT &next = curr_traj_.x.col(curr_t_);
                //joint_traj_.header.frame_id = "torso_lift_link";//goal->header.frame_id;
                joint_traj_.header.stamp = time_stamp_now;
                // NOTE: we assume that the message is already prepared
                //       and that joint_names are filled in for us
                for (int i = 0; i < joint_traj_.points[0].positions.size(); ++i) {
                        joint_traj_.points[0].positions[i] = next(i); // curr_joint_goal(i);
                        joint_traj_.points[0].velocities[i] = next(i + 7); // curr_joint_goal(i);
                }
                joint_traj_.points[0].time_from_start = ros::Duration(0.8f);
                
                // remember last goal
                last_goal_ = goal;
                last_goal_set_ = true;

                // remember last pos
                last_pos_ = center_;
                last_pos_set_ = true;

                // check if we found a goal fulfilling trajectory
                if (curr_traj_.x.col(curr_traj_.x.cols()-1).norm() < eps_) {
                        on_way_to_goal_ = true;
                } else {
                        on_way_to_goal_ = false;
                }
                
                // send action
                joint_pose_pub_.publish(joint_traj_);
                open_gripper();

                // TODO: send out marker array
                if (marker_pub_.getNumSubscribers() > 0) {
                        ROS_INFO("mpc sending markers");
                        visualization_msgs::MarkerArray markers;
                        VectorT tmp_pos(3);
                        int id = 0;
                        for (int p = 0; p < curr_traj_.x.cols(); ++p) {
                                r_arm_model_->fk(curr_traj_.x.col(p), tmp_pos);
                                visualization_msgs::Marker marker;
                                marker.header.frame_id = objects->pose.header.frame_id;
                                marker.header.stamp = ros::Time::now();
                                marker.type = marker.SPHERE;
                                marker.action = marker.ADD;
                                marker.id = id;
                                marker.scale.x = 0.02;
                                marker.scale.y = 0.02;
                                marker.scale.z = 0.02;
                                marker.color.a = 1.0;
                                marker.color.r = 1.0;
                                marker.color.g = 1.0;
                                marker.color.b = 0.0;
                                marker.pose.orientation.w = 1.0;
                                marker.pose.position.x = tmp_pos(0) + goal(0);
                                marker.pose.position.y = tmp_pos(1) + goal(1);
                                marker.pose.position.z = tmp_pos(2) + goal(2);
                                markers.markers.push_back(marker);
                                ++id;
                        }
                        marker_pub_.publish(markers);
                }
        }
};

int main(int argc, char *argv[]) {
        ros::init(argc, argv, "mpc_node");
        ros::NodeHandle n;
        ArmMPCJoints mpc_client(n);
        while (ros::ok()) {
                ros::spinOnce();
        }
}