Squashed 'modules/atom01_deploy/src/inference/' changes from f75ce36..56e7c8a

56e7c8a fix bug
a088d5b add policy mode switch log
6feffc5 fix bug
879864d fix policy changing logic
c47f1ca update conf

git-subtree-dir: modules/atom01_deploy/src/inference
git-subtree-split: 56e7c8a72be4bad25cf854a2a79b9744bb883145

Author: tcmofashi

config/inference_amp.yaml

@@ -52,4 +52,4 @@ inference_node:
              -3.14, 3.14,
              -3.14, 3.14,
              -3.14, 3.14]
-        gravity_z_upper: -0.7
+        gravity_z_upper: -0.3

config/inference_beyondmimic.yaml

@@ -25,9 +25,9 @@ inference_node:
         clip_actions: 100.0
         usd2urdf: [0, 6, 12, 1, 7, 13, 18, 2, 8, 14, 19, 3, 9, 15, 20, 4, 10, 16, 21, 5, 11, 17, 22]
         clip_cmd: 
-            [-0.0, 0.0,
-             -0.0, 0.0,
-             -0.0, 0.0]
+            [-0.4, 0.6,
+             -0.4, 0.4,
+             -0.8, 0.8]
         joint_default_angle: 
             [0.0, 0.0, -0.1, 0.3, -0.2, 0.0,
              0.0, 0.0, -0.1, 0.3, -0.2, 0.0, 0.0,

config/robot.yaml

@@ -23,7 +23,7 @@ motors:
     master_id_offset: 16
     motor_zero_offset: 
         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
-         0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+         0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.093,
          0.0, 0.0, 0.0, 0.0, 0.0,
          0.0 ,0.0, 0.0, 0.0, 0.0]
 

src/inference_node.cpp

@@ -56,6 +56,10 @@ void InferenceNode::setup_model(std::unique_ptr<ModelContext>& ctx, std::string
 
 void InferenceNode::reset() {
     is_running_.store(false);
+    is_interrupt_.store(false);
+    is_beyondmimic_.store(false);
+    obs_.resize(obs_num_);
+    active_ctx_ = normal_ctx_.get();
     std::fill(obs_.begin(), obs_.end(), 0.0f);
     std::fill(joint_pos_.begin(), joint_pos_.end(), 0.0f);
     std::fill(joint_vel_.begin(), joint_vel_.end(), 0.0f);
@@ -64,17 +68,19 @@ void InferenceNode::reset() {
     std::fill(cmd_vel_.begin(), cmd_vel_.end(), 0.0f);
     std::fill(quat_.begin(), quat_.end(), 0.0f);
     std::fill(ang_vel_.begin(), ang_vel_.end(), 0.0f);
-    if (active_ctx_) {
-        std::fill(active_ctx_->input_buffer.begin(), active_ctx_->input_buffer.end(), 0.0f);
-        std::fill(active_ctx_->output_buffer.begin(), active_ctx_->output_buffer.end(), 0.0f);
+    if (normal_ctx_) {
+        std::fill(normal_ctx_->input_buffer.begin(), normal_ctx_->input_buffer.end(), 0.0f);
+        std::fill(normal_ctx_->output_buffer.begin(), normal_ctx_->output_buffer.end(), 0.0f);
+    }
+    if (motion_ctx_) {
+        std::fill(motion_ctx_->input_buffer.begin(), motion_ctx_->input_buffer.end(), 0.0f);
+        std::fill(motion_ctx_->output_buffer.begin(), motion_ctx_->output_buffer.end(), 0.0f);
     }
     std::fill(act_.begin(), act_.end(), 0.0f);
     std::fill(last_act_.begin(), last_act_.end(), 0.0f);
     std::fill(joint_torques_.begin(), joint_torques_.end(), 0.0f);
     is_first_frame_ = true;
     motion_frame_ = 0;
-    is_interrupt_.store(false);
-    is_beyondmimic_.store(false);
     if(use_interrupt_){
         std::fill(interrupt_action_.begin(), interrupt_action_.end(), 0.0f);
     }
@@ -277,6 +283,9 @@ int main(int argc, char **argv) {
         RCLCPP_INFO(node->get_logger(), "Press 'X' to reset motors");
         RCLCPP_INFO(node->get_logger(), "Press 'B' to start/pause inference");
         RCLCPP_INFO(node->get_logger(), "Press 'Y' to switch between joystick and /cmd_vel control");
+        if (node->use_interrupt_ || node->use_beyondmimic_){
+            RCLCPP_INFO(node->get_logger(), "Press 'LB' to switch policy mode");
+        }
         executor.spin();
     } catch (const std::exception &e) {
         RCLCPP_FATAL(rclcpp::get_logger("main"), "Exception caught: %s", e.what());

src/inference_node.hpp

@@ -146,12 +146,12 @@ class InferenceNode : public rclcpp::Node {
         size_t num_inputs;
         size_t num_outputs;
     };
+    bool use_interrupt_, use_beyondmimic_, use_attn_enc_;
    private:
     std::shared_ptr<RobotInterface> robot_;
     int offline_threshold_ = 10;
     std::atomic<bool> is_running_{false}, is_joy_control_{true}, is_interrupt_{false}, is_beyondmimic_{false};
     std::string model_name_, model_path_, motion_name_, motion_path_, motion_model_name_, motion_model_path_, perception_obs_topic_;
-    bool use_interrupt_, use_beyondmimic_, use_attn_enc_;
     int obs_num_, motion_obs_num_, perception_obs_num_, frame_stack_, motion_frame_stack_, joint_num_;
     int decimation_;
     std::unique_ptr<Ort::Env> env_;

src/robot_interface.cpp

@@ -156,7 +156,7 @@ void RobotInterface::reset_joints(std::vector<double> joint_default_angle) {
     }
 
     exec_motors_parallel([this, &joint_default_angle](std::shared_ptr<MotorDriver>& motor, int idx) {
-        motor->motor_mit_cmd(joint_default_angle[idx] * robot_cfg_->motor_sign_[idx], 0.0f, robot_cfg_->kp_[idx]/2.0f, robot_cfg_->kd_[idx]/2.0f, 0.0f);
+        motor->motor_mit_cmd(joint_default_angle[idx] * robot_cfg_->motor_sign_[idx], 0.0f, robot_cfg_->kp_[idx]/2.5f, robot_cfg_->kd_[idx], 0.0f);
     });
     std::this_thread::sleep_for(std::chrono::milliseconds(1000));
     exec_motors_parallel([this, &joint_default_angle](std::shared_ptr<MotorDriver>& motor, int idx) {

src/ros_interface.cpp

@@ -32,6 +32,7 @@ void InferenceNode::load_config() {
     this->declare_parameter<std::vector<double>>("clip_cmd", std::vector<double>{});
     this->declare_parameter<std::vector<double>>("joint_default_angle", std::vector<double>{});
     this->declare_parameter<std::vector<double>>("joint_limits", std::vector<double>{});
+    this->declare_parameter<float>("gravity_z_upper", -0.5);
 
 
     this->get_parameter("model_name", model_name_);
@@ -65,6 +66,7 @@ void InferenceNode::load_config() {
     this->get_parameter("clip_cmd", clip_cmd_);
     this->get_parameter("joint_default_angle", joint_default_angle_);
     this->get_parameter("joint_limits", joint_limits_);
+    this->get_parameter("gravity_z_upper", gravity_z_upper_);
 
 
     model_path_ = std::string(ROOT_DIR) + "models/" + model_name_;
@@ -96,6 +98,7 @@ void InferenceNode::load_config() {
     print_vector<double>("clip_cmd", clip_cmd_);
     print_vector<double>("joint_default_angle", joint_default_angle_);
     print_vector<double>("joint_limits", joint_limits_);
+    RCLCPP_INFO(this->get_logger(), "gravity_z_upper: %f", gravity_z_upper_);
 }
 
 void InferenceNode::subs_joy_callback(const std::shared_ptr<sensor_msgs::msg::Joy> msg) {
@@ -147,11 +150,16 @@ void InferenceNode::subs_joy_callback(const std::shared_ptr<sensor_msgs::msg::Jo
     if (use_interrupt_ || use_beyondmimic_) {
         if (msg->buttons[4] == 1 && msg->buttons[4] != last_button4_) {
             if(use_interrupt_){
+                is_running_.store(false);
+                RCLCPP_INFO(this->get_logger(), "Inference paused");
+                std::this_thread::sleep_for(std::chrono::milliseconds(50));
                 is_interrupt_.store(!is_interrupt_.load());
+                std::this_thread::sleep_for(std::chrono::milliseconds(50));
                 RCLCPP_INFO(this->get_logger(), "Interrupt mode %s", is_interrupt_.load() ? "enabled" : "disabled");
             }
-            if(use_beyondmimic_){
+            else if(use_beyondmimic_){
                 is_running_.store(false);
+                RCLCPP_INFO(this->get_logger(), "Inference paused");
                 std::this_thread::sleep_for(std::chrono::milliseconds(50));
                 is_beyondmimic_.store(!is_beyondmimic_.load());
                 bool is_beyondmimic = is_beyondmimic_.load();
@@ -170,10 +178,10 @@ void InferenceNode::subs_joy_callback(const std::shared_ptr<sensor_msgs::msg::Jo
                 std::fill(active_ctx_->output_buffer.begin(), active_ctx_->output_buffer.end(), 0.0f);
                 std::fill(act_.begin(), act_.end(), 0.0f);
                 std::fill(last_act_.begin(), last_act_.end(), 0.0f);
+                std::fill(joint_torques_.begin(), joint_torques_.end(), 0.0f);
                 is_first_frame_ = true;
                 motion_frame_ = 0;
                 std::this_thread::sleep_for(std::chrono::milliseconds(50));
-                is_running_.store(true);
                 RCLCPP_INFO(this->get_logger(), "Beyondmimic mode %s", is_beyondmimic ? "enabled" : "disabled");
             }
         }