initial commit

This commit is contained in:
huangweifeng
2026-10-22 09:23:54 +08:00
commit 63f557a39f
4448 changed files with 584487 additions and 0 deletions
+179
View File
@@ -0,0 +1,179 @@
#pragma once
#include "policy_runner_base.hpp"
#include <array>
#include <cmath>
#include <onnxruntime_cxx_api.h>
#include <stdexcept>
/** Stage-2 FireAnt policy adapter (63 observations, 8 actions, 50 Hz). */
class FireAntPolicyRunner : public PolicyRunnerBase {
private:
static constexpr int kMotorNum = 8;
static constexpr int kObservationDim = 63;
static constexpr int kActionDim = 8;
static constexpr float kWheelRadius = 0.096f;
// Policy: FL/FR/HL/HR knees, then FL/FR/HL/HR wheels.
// Robot/DDS: FL knee/wheel, FR knee/wheel, HL knee/wheel, HR knee/wheel.
const std::array<int, kActionDim> policy_to_robot_{{0, 2, 4, 6, 1, 3, 5, 7}};
const std::array<float, 4> knee_default_{{0.5f, 0.5f, -0.5f, -0.5f}};
std::string policy_path_;
VecXf observation_ = VecXf::Zero(kObservationDim);
VecXf last_action_ = VecXf::Zero(kActionDim);
RobotAction robot_action_;
Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "FireAntPolicyRunner"};
Ort::SessionOptions session_options_;
Ort::Session session_{nullptr};
Ort::MemoryInfo memory_info_{Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)};
const char *input_names_[1] = {"obs"};
const char *output_names_[1] = {"actions"};
const std::array<int64_t, 2> input_shape_{{1, kObservationDim}};
static float Clamp(float value, float low, float high) {
return std::max(low, std::min(value, high));
}
void ValidateModel() {
if (session_.GetInputCount() != 1 || session_.GetOutputCount() != 1) {
throw std::runtime_error("FireAnt ONNX must have exactly one input and one output");
}
auto input_dims = session_.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
auto output_dims = session_.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
if (input_dims.size() != 2 || input_dims[1] != kObservationDim) {
throw std::runtime_error("FireAnt ONNX input must be [batch, 63]");
}
if (output_dims.size() != 2 || output_dims[1] != kActionDim) {
throw std::runtime_error("FireAnt ONNX output must be [batch, 8]");
}
}
VecXf Infer() {
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
memory_info_, observation_.data(), observation_.size(), input_shape_.data(), input_shape_.size());
auto outputs = session_.Run(
Ort::RunOptions{nullptr}, input_names_, &input_tensor, 1, output_names_, 1);
const float *data = outputs[0].GetTensorData<float>();
return Eigen::Map<const VecXf>(data, kActionDim);
}
public:
FireAntPolicyRunner(const std::string &policy_name, const std::string &policy_path)
: PolicyRunnerBase(policy_name), policy_path_(policy_path) {
if (access(policy_path_.c_str(), F_OK) != 0) {
throw std::runtime_error("FireAnt policy not found: " + policy_path_);
}
SetDecimation(4); // 200 Hz state loop / 4 = the 50 Hz training policy rate.
session_options_.SetIntraOpNumThreads(2);
session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
session_ = Ort::Session(env_, policy_path_.c_str(), session_options_);
ValidateModel();
robot_action_.goal_joint_pos = VecXf::Zero(kMotorNum);
robot_action_.goal_joint_vel = VecXf::Zero(kMotorNum);
robot_action_.kp = VecXf::Zero(kMotorNum);
robot_action_.kd = VecXf::Zero(kMotorNum);
robot_action_.tau_ff = VecXf::Zero(kMotorNum);
for (int leg = 0; leg < 4; ++leg) {
robot_action_.kp(2 * leg) = 80.0f;
robot_action_.kd(2 * leg) = 2.0f;
robot_action_.kd(2 * leg + 1) = 0.6f;
}
}
void DisplayPolicyInfo() override {
std::cout << "[FireAnt] ONNX: " << policy_path_ << "\n"
<< "[FireAnt] observations=63, actions=8, policy_rate=50 Hz" << std::endl;
}
void OnEnter() override {
run_cnt_ = 0;
cmd_vel_input_.setZero();
last_action_.setZero();
}
RobotAction getRobotAction(const RobotBasicState &ro, const UserCommand &uc) override {
if (ro.joint_pos.size() != kMotorNum || ro.joint_vel.size() != kMotorNum ||
ro.joint_tau.size() != kMotorNum) {
throw std::runtime_error("FireAnt robot state must contain 8 joints");
}
// Keyboard commands act as a deadman switch. Stage-2 sampled only positive
// forward commands, so zero-command inference can otherwise cause creeping.
if (std::abs(uc.forward_vel_scale) + std::abs(uc.side_vel_scale) +
std::abs(uc.turnning_vel_scale) < 0.05f) {
last_action_.setZero();
robot_action_.goal_joint_pos.setZero();
robot_action_.goal_joint_vel.setZero();
for (int leg = 0; leg < 4; ++leg) {
robot_action_.goal_joint_pos(2 * leg) = knee_default_[leg];
}
++run_cnt_;
return robot_action_;
}
int cursor = 0;
observation_.segment<3>(cursor) = 0.25f * ro.base_omega;
cursor += 3;
observation_.segment<3>(cursor) = ro.base_rot_mat.transpose() * Vec3f(0.0f, 0.0f, -1.0f);
cursor += 3;
observation_.segment<3>(cursor) = Vec3f(
uc.forward_vel_scale, uc.side_vel_scale, uc.turnning_vel_scale);
cursor += 3;
// Relative joint position in policy order. Continuous wheel positions are intentionally zero.
for (int leg = 0; leg < 4; ++leg) {
observation_(cursor + leg) = ro.joint_pos(policy_to_robot_[leg]) - knee_default_[leg];
observation_(cursor + 4 + leg) = 0.0f;
}
cursor += 8;
for (int i = 0; i < kActionDim; ++i) {
observation_(cursor + i) = 0.05f * ro.joint_vel(policy_to_robot_[i]);
}
cursor += 8;
observation_.segment(cursor, kActionDim) = last_action_;
cursor += 8;
VecXf contact = ro.contact_force.size() == 9 ? ro.contact_force : VecXf::Zero(9);
for (int i = 0; i < 4; ++i) observation_(cursor + i) = contact(i) > 2.0f ? 1.0f : 0.0f;
cursor += 4;
for (int i = 0; i < 5; ++i) observation_(cursor + i) = contact(4 + i) > 2.0f ? 1.0f : 0.0f;
cursor += 5;
for (int i = 0; i < 4; ++i) observation_(cursor + i) = Clamp(contact(i) / 100.0f, 0.0f, 2.0f);
cursor += 4;
for (int i = 0; i < 5; ++i) observation_(cursor + i) = Clamp(contact(4 + i) / 100.0f, 0.0f, 2.0f);
cursor += 5;
const float forward_speed = std::abs(ro.base_lin_vel(0));
for (int leg = 0; leg < 4; ++leg) {
const float wheel_speed = std::abs(ro.joint_vel(2 * leg + 1) * kWheelRadius);
observation_(cursor + leg) = 0.5f * Clamp(std::abs(wheel_speed - forward_speed), 0.0f, 3.0f);
}
cursor += 4;
for (int i = 0; i < kActionDim; ++i) {
observation_(cursor + i) = Clamp(ro.joint_tau(policy_to_robot_[i]) / 80.0f, -2.0f, 2.0f);
}
cursor += 8;
if (cursor != kObservationDim || !observation_.allFinite()) {
throw std::runtime_error("Invalid FireAnt policy observation");
}
VecXf action = Infer();
if (!action.allFinite()) throw std::runtime_error("FireAnt policy produced non-finite actions");
for (int i = 0; i < kActionDim; ++i) action(i) = Clamp(action(i), -100.0f, 100.0f);
last_action_ = action;
robot_action_.goal_joint_pos.setZero();
robot_action_.goal_joint_vel.setZero();
for (int leg = 0; leg < 4; ++leg) {
robot_action_.goal_joint_pos(2 * leg) = knee_default_[leg] + 0.25f * action(leg);
robot_action_.goal_joint_vel(2 * leg + 1) = 5.0f * action(4 + leg);
}
++run_cnt_;
return robot_action_;
}
};
+255
View File
@@ -0,0 +1,255 @@
/**
* @file m20_policy_runner.hpp
* @brief m20_policy_runner
* @author Bo (Percy) Peng
* @version 1.0
* @date 2025-11-07
*
* @copyright Copyright (c) 2025 DeepRobotics
*
*/
#pragma once
#define PI 3.14159265358979323846
#include "policy_runner_base.hpp"
#include <ctime>
#include <cmath>
#include <utility>
#include <onnxruntime_cxx_api.h>
#include <onnxruntime_c_api.h>
class M20PolicyRunner : public PolicyRunnerBase {
private:
VecXf kp_, kd_;
VecXf dof_default_eigen_policy, dof_default_eigen_robot;
Vec3f max_cmd_vel_, gravity_direction = Vec3f(0., 0., -1.);
VecXf dof_pos_default_;
timespec system_time;
const int motor_num = 16;
const int observation_dim = 57;
const int action_dim = 16;
float agent_timestep = 0.02;
float current_time;
bool is_fallen = true;
VecXf joint_pos_rl = VecXf(action_dim);// in rl squenece
VecXf joint_vel_rl = VecXf(action_dim);
const std::string policy_path_;
float omega_scale_ = 0.25;
float dof_vel_scale_ = 0.05;
VecXf imu_w_eigen, base_acc_eigen, motor_p_eigen, motor_v_eigen,
current_action_eigen, last_action_eigen, current_observation_, projected_gravity,
tmp_action_eigen;
RobotAction robot_action;
std::vector<std::string> robot_order = {
"fl_hipx_joint", "fl_hipy_joint", "fl_knee_joint", "fl_wheel_joint",
"fr_hipx_joint", "fr_hipy_joint", "fr_knee_joint", "fr_wheel_joint",
"hl_hipx_joint", "hl_hipy_joint", "hl_knee_joint", "hl_wheel_joint",
"hr_hipx_joint", "hr_hipy_joint", "hr_knee_joint", "hr_wheel_joint"};
std::vector<std::string> policy_order = {
"fl_hipx_joint", "fl_hipy_joint", "fl_knee_joint",
"fr_hipx_joint", "fr_hipy_joint", "fr_knee_joint",
"hl_hipx_joint", "hl_hipy_joint", "hl_knee_joint",
"hr_hipx_joint", "hr_hipy_joint", "hr_knee_joint",
"fl_wheel_joint", "fr_wheel_joint", "hl_wheel_joint", "hr_wheel_joint",
};
std::vector<float> action_scale_robot = {0.125, 0.25, 0.25, 5,
0.125, 0.25, 0.25, 5,
0.125, 0.25, 0.25, 5,
0.125, 0.25, 0.25, 5};
Ort::SessionOptions session_options_;
Ort::Session session_{nullptr};
Ort::Env env_;
std::vector<int> robot2policy_idx, policy2robot_idx;
const char* input_names_[1] = {"obs"}; // must keep the same as model export
const char* output_names_[1] = {"actions"};
VecXf command;
Ort::MemoryInfo memory_info{nullptr};
const std::array<int64_t, 2> input_observationShape = {1, observation_dim};
float time_step = 0.;
int stop_count = 1000;
public:
M20PolicyRunner(const std::string &policy_name, const std::string &policy_path) :
PolicyRunnerBase(policy_name), policy_path_(policy_path),env_(ORT_LOGGING_LEVEL_WARNING, "M20PolicyRunner"),
session_options_{},
session_{nullptr},
memory_info(Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)) {
dof_default_eigen_policy.setZero(action_dim);
dof_default_eigen_robot.setZero(action_dim);
dof_default_eigen_policy << 0.0, -0.6, 1.0,
0.0, -0.6, 1.0,
0.0, 0.6, -1.0,
0.0, 0.6, -1.0,
0.0, 0.0, 0.0, 0.0;
dof_default_eigen_robot << 0.0, -0.6, 1.0, 0.0,
0.0, -0.6, 1.0, 0.0,
0.0, 0.6, -1.0, 0.0,
0.0, 0.6, -1.0, 0.0;
SetDecimation(4);
session_options_.SetIntraOpNumThreads(4);
session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
if (access(policy_path_.c_str(), F_OK) != 0) {
std::cerr << "Model file not found: " << policy_path_ << std::endl;
throw std::runtime_error("Model file missing");
}
// 加载模型
session_ = Ort::Session(env_, policy_path_.c_str(), session_options_);
kp_ = Vec4f(80, 80, 80, 0.).replicate(4, 1);
kd_ = Vec4f(2, 2, 2, 0.6).replicate(4, 1);
robot2policy_idx = generate_permutation(robot_order, policy_order);
policy2robot_idx = generate_permutation(policy_order, robot_order);
// for (int i = 0; i < action_dim; ++i){
// std::cout << "robot2policy_idx[" << i << "]: " << robot2policy_idx[i] << std::endl;
// std::cout << "policy2robot_idx[" << i << "]: " << policy2robot_idx[i] << std::endl;
// }
robot_action.kp = kp_;
robot_action.kd = kd_;
robot_action.tau_ff = VecXf::Zero(motor_num);
robot_action.goal_joint_pos = VecXf::Zero(motor_num);
robot_action.goal_joint_vel = VecXf::Zero(motor_num);
current_observation_.setZero(observation_dim);
last_action_eigen.setZero(action_dim);
tmp_action_eigen.setZero(action_dim);
current_action_eigen.setZero(action_dim);
memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
}
~M20PolicyRunner() override = default;
std::vector<int> generate_permutation(
const std::vector<std::string>& from,
const std::vector<std::string>& to,
int default_index = 0)
{
std::unordered_map<std::string, int> idx_map;
for (int i = 0; i < from.size(); ++i) {
idx_map[from[i]] = i;
}
std::vector<int> perm;
for (const auto& name : to) {
auto it = idx_map.find(name);
if (it != idx_map.end()) {
perm.push_back(it->second);
} else {
perm.push_back(default_index); // 如果找不到,就填默认值
}
}
return perm;
}
void DisplayPolicyInfo(){}
void OnEnter() {
run_cnt_ = 0;
cmd_vel_input_.setZero();
last_action_eigen.setZero(action_dim);
tmp_action_eigen.setZero(action_dim);
motor_p_eigen.setZero(12);
motor_v_eigen.setZero(motor_num);
}
VecXf Onnx_infer(VecXf current_observation){
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
memory_info,
current_observation.data(),
current_observation.size(),
input_observationShape.data(),
input_observationShape.size()
);
std::vector<Ort::Value> inputs;
inputs.emplace_back(std::move(input_tensor)); // 避免拷贝构造
auto outputs = session_.Run(
Ort::RunOptions{nullptr},
input_names_,
inputs.data(),
1,
output_names_,
1
);
float* action_data = outputs[0].GetTensorMutableData<float>();
Eigen::Map<Eigen::VectorXf> action_map(action_data, action_dim);
return VecXf(action_map); // 返回一个Eigen向量的副本
}
RobotAction getRobotAction(const RobotBasicState &ro, const UserCommand &uc) {
Vec3f base_omgea = ro.base_omega * omega_scale_;
Vec3f projected_gravity = ro.base_rot_mat.inverse() * gravity_direction;
Vec3f command = Vec3f(uc.forward_vel_scale, uc.side_vel_scale, uc.turnning_vel_scale);
for (int i = 0; i < action_dim; ++i){
joint_pos_rl(i) = ro.joint_pos(robot2policy_idx[i]);
joint_vel_rl(i) = ro.joint_vel(robot2policy_idx[i]) * dof_vel_scale_;
}
joint_pos_rl.segment(12, 4).setZero();
joint_pos_rl -= dof_default_eigen_policy;
current_observation_<<base_omgea,
projected_gravity,
command,
joint_pos_rl,
joint_vel_rl,
last_action_eigen;
current_action_eigen = Onnx_infer(current_observation_);
last_action_eigen = current_action_eigen;
for (int i = 0; i < action_dim; ++i){
tmp_action_eigen(i) = current_action_eigen(policy2robot_idx[i]);
tmp_action_eigen(i) *= action_scale_robot[i];
}
tmp_action_eigen += dof_default_eigen_robot;
for (int i = 0; i < 4; ++i){
robot_action.goal_joint_pos.segment(i*4, 3) = tmp_action_eigen.segment(i*4, 3);
robot_action.goal_joint_vel(i*4+3) = tmp_action_eigen(i*4+3);
}
++run_cnt_;
++time_step;
return robot_action;
}
void setDefaultJointPos(const VecXf& pos){
dof_pos_default_.setZero(motor_num);
for(int i=0;i<motor_num;++i) {
dof_pos_default_(i) = pos(i);
}
}
double getCurrentTime() {
clock_gettime(1, &system_time);
return system_time.tv_sec + system_time.tv_nsec / 1e9;
}
};
+58
View File
@@ -0,0 +1,58 @@
/**
* @file policy_runner_base.hpp
* @brief policy runner base class
* @author DeepRobotics
* @version 1.0
* @date 2025-11-07
*
* @copyright Copyright (c) 2025 DeepRobotics
*
*/
#pragma once
#include "common_types.h"
#include "basic_function.hpp"
#include "onnxruntime_cxx_api.h"
using namespace types;
class PolicyRunnerBase{
public:
PolicyRunnerBase(std::string policy_name):policy_name_(policy_name){
vel_delta_const_ << 0.3, 0.2, 0.3;
cmd_vel_input_.setZero();
}
virtual ~PolicyRunnerBase(){}
/**
* @brief use to display policy detail
*/
virtual void DisplayPolicyInfo(){
}
/**
* @brief Get the robot action object by run your policy
* @return RobotAction
*/
virtual RobotAction getRobotAction(const RobotBasicState&, const UserCommand&) = 0;
/**
* @brief execute function when first entering policy runner
*/
virtual void OnEnter() = 0;
/**
* @brief Set the decimation
* @param d decimation
*/
virtual void SetDecimation(int d){
decimation_ = d;
}
const std::string policy_name_;
int decimation_;
int run_cnt_;
Vec3f vel_delta_const_, cmd_vel_input_;
};