initial commit
This commit is contained in:
huangweifeng
@@ -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_;
|
||||
}
|
||||
};
|
||||
Reference in New Issue
Block a user