3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Sim-to-Real実装:PyTorchで学習した強化学習モデルをONNX経由でSTM32にデプロイする

3
Posted at

3軸リアクションホイール倒立振子を作るアドカレの11日目です。

1. はじめに:Sim-to-Realの架け橋

これまでの工程で、シミュレーションとハードウェア設計の両輪が揃いました。

  • Day 7: PyTorch (Stable Baselines 3) を用いたPPOアルゴリズムによる姿勢制御モデルの学習。
  • Day 9-10: STM32G474を搭載した制御基板の回路設計。

本日は、Python環境で学習させた「脳(AIモデル)」を、C言語ベースの組み込みマイコン (STM32) に移植する Sim-to-Real(シミュレーションから実環境への移行) の工程を行います。
具体的には、PyTorchモデルを中間表現である ONNX 形式に変換し、STMicroelectronics公式ツール X-CUBE-AI を用いて最適化されたCコードを生成・実装します。

本日の実装目標

  1. 学習済みモデル(.zip)をONNX形式に変換する。
  2. STM32G474上で動作可能なCソースコードを自動生成する。
  3. 実機上で推論を実行し、トルク指令値がUART経由で出力されることを確認する。

2. Step 1: ONNX形式への変換 (Python)

まずは、Stable Baselines 3 (SB3) で学習したモデルを、ディープラーニングモデルの共通フォーマットである ONNX (Open Neural Network Exchange) に変換します。

SB3モデルをONNX化する際の壁

通常、SB3の model.predict() メソッドは内部で観測値の前処理を行い、結果をNumPy配列として返します。しかし、変換に用いる torch.onnx.export は PyTorch Tensor の計算グラフ(演算の流れ)を追跡(トレース)する仕組みであるため、途中でNumPy配列に変換されるとトレースが途切れ、変換エラーとなってしまいます。

そこで、nn.Module を継承したラッパークラスを作成し、Policy Networkの内部レイヤー(extract_features -> mlp_extractor -> action_net)を直接呼び出すことで、Tensor型を維持したまま計算グラフを構築する手法をとります。

変換における「ダミー入力」の重要性

以下のスクリプト内では dummy_input として (1, 6) のランダムな数値を生成していますが、これには重要な意味があります。

ONNXへの変換は、モデルの設計図を静的に読むのではなく、「実際にデータを1回流してみて、計算がどのような順序で行われるかを記録(トレース)する」 という手法で行われます。そのため、値自体はランダムで構いませんが、「実際に計算を通すことができるデータ」 を渡す必要があります。

また、ここで指定するデータの形(Shape)が、そのまま生成されるONNXモデルの仕様として固定されます。
今回の 6 という次元数は、Day 7 の学習フェーズで使用したパラメータ(Observation Space)と完全に一致させる必要があります。学習済みモデルはこの6つの入力を受け取る前提で訓練されているため、変換時にも全く同じ次元数のデータ型紙を渡さなければ、整合性が取れず変換に失敗します。

【入力次元数「6」の内訳】
今回AIに入力する6つのデータは、「現在のCubliの状態」 を表す数値です。

インデックス 項目名 単位 説明
0 Roll Angle [rad] X軸周りの傾き(ロール角)
1 Pitch Angle [rad] Y軸周りの傾き(ピッチ角)
2 Gyro X [rad/s] X軸周りの角速度
3 Gyro Y [rad/s] Y軸周りの角速度
4 Gyro Z [rad/s] Z軸周りの角速度
5 Wheel Speed [rad/s] ホイール回転速度(飽和防止用)

変換スクリプト: export_onnx.py

import torch
import torch.nn as nn
from stable_baselines3 import PPO

# 1. 学習済みモデルのロード
# device="cpu"を指定することで、CPU環境での推論前提のモデルとしてロードします
model_path = "models/PPO_Robust/final_model.zip"
model = PPO.load(model_path, device="cpu")

# 2. STM32実装用ラッパーモデルの定義
# 推論に必要なActorネットワークのみを抽出して実行するクラス
class OnnxablePolicy(nn.Module):
    def __init__(self, policy):
        super().__init__()
        self.policy = policy

    def forward(self, observation):
        # NOTE: SB3の内部メソッドを直接呼ぶことでTensorのまま処理を流す
        
        # 1. 特徴量抽出 (Preprocessing)
        features = self.policy.extract_features(observation)
        # 2. 潜在特徴の算出 (Latent features)
        latent_pi, _ = self.policy.mlp_extractor(features)
        # 3. 行動(トルク)の算出 (Deterministic Action)
        action = self.policy.action_net(latent_pi)
        
        return action

# ラッパーの適用
onnx_policy = OnnxablePolicy(model.policy)

# 3. ダミー入力の生成 (入力次元数: 6)
# 想定入力: [Roll_err, Pitch_err, Yaw_err, GyroX, GyroY, GyroZ]
dummy_input = torch.randn(1, 6)

# 4. ONNXエクスポート実行
torch.onnx.export(
    onnx_policy,
    dummy_input,
    "cubli_policy.onnx",    # 出力ファイル名
    opset_version=11,       # X-CUBE-AI推奨バージョン
    input_names=["input"],
    output_names=["output"]
)

print("Conversion Complete: cubli_policy.onnx generated.")

このスクリプトを実行すると、cubli_policy.onnx が生成されます。これがSTM32に渡す「脳のデータ」となります。

3.1 X-CUBE-AIの導入と設定

  1. STM32CubeMXを開き、[Software Packs] -> [Select Components] から STMicroelectronics.X-CUBE-AI を選択し、Coreコンポーネントを有効化します。
  2. 左側のメニューに追加された [X-CUBE-AI] タブを開き、[Add network] をクリック。
  3. Model inputs に先ほど生成した cubli_policy.onnx を指定します。
    image.png

3.2 リソース使用量の解析 (Analyze)

ターゲットデバイス (STM32G474RET6) に対するメモリ使用量と計算負荷を解析します。[Analyze] ボタンをクリックすると、詳細なレポートが表示されます。
image.png

【解析ログの詳細評価】
実際の解析ログから、このAIモデルがマイコンに与える負荷を読み解きます。

Model file:      cubli_policy.onnx
Total Flash:     28434 B (27.77 KiB)
    Weights:     19212 B (18.76 KiB)
    Library:     9222 B (9.01 KiB)
Total Ram:       2348 B (2.29 KiB)
    Activations: 524 B
    Library:     1824 B (1.78 KiB)
MACC:            6,083
  • Flash使用量: 27.77 KiB (STM32G474 512KB の 約5.4%)
    • 内訳の大部分(約19KB)は学習済みの重みパラメータ(Weights)です。残りは推論エンジンのプログラムコードです。容量には十分な余裕があります。
  • RAM使用量: 2.29 KiB (STM32G474 128KB の 約1.8%)
    • 作業領域(Activations)はわずか524バイトです。これはモデルが小さく、中間層のニューロン数が少ないためです。他の制御用変数やスタック領域を圧迫する心配はありません。
  • MACC (積和演算数): 6,083回
    • これは推論1回あたりの計算量を示します。STM32G4 (170MHz) の処理能力を考慮すると、推論にかかる時間は 数マイクロ秒〜数十マイクロ秒 程度と推測されます。
    • 今回の制御周期(1kHz = 1000µs)に対して、AI処理時間は誤差レベルの短さであり、リアルタイム制御への影響は皆無です。

確認後、[Generate Code] を実行し、Cソースコードおよびヘッダファイルを生成します。

4. Step 3: 推論エンジンの実装 (STM32CubeIDE)

生成されたコードは Middlewares/ST/AI ディレクトリ等に配置されます。これらを呼び出すユーザコードを main.c に実装します。

4.1 ヘッダインクルードとバッファ確保

AIエンジンのヘッダを読み込み、モデルの重みや入出力データを格納するバッファを確保します。ここで定義する in_data のサイズが、先ほどのダミー入力のサイズ(6)と対応しています。
AI_ALIGNED(32) マクロを使用し、メモリアライメントを調整することが重要です。

/* USER CODE BEGIN Includes */
#include "ai_platform.h"
#include "network.h"     // 生成されたモデル定義 (network.hという名前は設定で可変)
#include "network_data.h"
#include <stdio.h>       // printf用
/* USER CODE END Includes */

/* USER CODE BEGIN PV */
// AIハンドラ・アクティベーションバッファ(中間層の演算結果保持用)
ai_handle network = AI_HANDLE_NULL;
AI_ALIGNED(32) ai_u8 activations[AI_NETWORK_DATA_ACTIVATIONS_SIZE];

// 入出力データバッファ
// In: [Roll, Pitch, GyroX, GyroY, GyroZ, WheelSpeed] -> 6次元
AI_ALIGNED(32) ai_float in_data[AI_NETWORK_IN_1_SIZE];
// Out: [Torque_X, Torque_Y, Torque_Z] -> 3次元
AI_ALIGNED(32) ai_float out_data[AI_NETWORK_OUT_1_SIZE];

// バッファポインタ構造体
ai_buffer *ai_input;
ai_buffer *ai_output;
/* USER CODE END PV */

4.2 推論エンジンの初期化

main 関数の初期化フェーズ(whileループの前)で、ニューラルネットワークのインスタンスを作成します。

  /* USER CODE BEGIN 2 */
  printf("AI Initialization...\r\n");

  // 1. ネットワークインスタンスの作成
  ai_error err;
  err = ai_network_create(&network, AI_NETWORK_DATA_CONFIG);
  if (err.type != AI_ERROR_NONE) {
      printf("AI Create Error: %d\r\n", err.type);
      Error_Handler();
  }

  // 2. 重みパラメータとワークバッファの割り当て
  ai_network_params params = AI_NETWORK_PARAMS_INIT(
      AI_NETWORK_DATA_WEIGHTS(ai_network_data_weights_get()),
      AI_NETWORK_DATA_ACTIVATIONS(activations)
  );

  if (!ai_network_init(network, &params)) {
      printf("AI Init Error\r\n");
      Error_Handler();
  }

  // 3. 入出力ポインタの取得と実データバッファの紐付け
  ai_input = ai_network_inputs_get(network, NULL);
  ai_output = ai_network_outputs_get(network, NULL);
  
  // 確保しておいた配列のアドレスをポインタにセット
  ai_input[0].data = AI_HANDLE_PTR(in_data);
  ai_output[0].data = AI_HANDLE_PTR(out_data);

  printf("AI Ready!\r\n");
  /* USER CODE END 2 */

4.3 推論実行ループ

現状はセンサが接続されていないため、ダミーデータを入力して推論が正常に回るかを確認します。

  /* USER CODE BEGIN 3 */
  while (1) {
    // 1. センサ入力 (テスト用ダミー値)
    // 本番環境ではここに IMU からのセンサ値を正規化して代入する
    for (int i = 0; i < AI_NETWORK_IN_1_SIZE; i++) {
        in_data[i] = 0.1f * (float)(i + 1); 
    }

    // 2. 推論実行 (Inference)
    ai_i32 batch = ai_network_run(network, ai_input, ai_output);

    if (batch != 1) {
        printf("Inference Error\r\n");
    } else {
        // 3. 結果出力 (モータトルク指令値)
        // 出力範囲は学習環境(PyTorch)のTanh出力なので -1.0 ~ 1.0 になるはず
        printf("In:[Dummy] -> Out(Torque): %.3f, %.3f, %.3f\r\n", 
                out_data[0], out_data[1], out_data[2]);
    }
    
    HAL_Delay(500); 
  }
  /* USER CODE END 3 */

5. 実装時のトラブルシューティング (Tips)

開発中に遭遇した技術的な問題とその解決策を共有します。

  1. 浮動小数点のprintf出力不可:
    printf%f を使用しても値が表示されない現象は、STM32CubeIDEのデフォルト設定(Newlib-nano)で浮動小数点サポートが無効化されているため発生します。[Project Properties] -> [C/C++ Build] -> [Settings] -> [MCU Settings] にて、Use float with printf from newlib-nano (-u _printf_float) にチェックを入れることで解決します。

  2. UART出力不可:
    プログラムは動作しているのにTeraTermに何も表示されない場合、ターゲットボードのVCP (Virtual COM Port) 配線とコード上のUARTインスタンスが不一致である可能性があります。syscalls.c 内の _write 関数で使用しているUARTハンドル(例: &huart1)が、実際の回路と合っているか確認が必要です。

6. 実行結果とまとめ

実機へ書き込みを行い、UARTログを確認しました。

AI Initialization...
AI Ready!
In:[Dummy] -> Out(Torque): -0.428, -0.291, 0.668
In:[Dummy] -> Out(Torque): -0.428, -0.291, 0.668
...

結果:
推論エンジンが正常に初期化され、エラーなく実行されていることを確認しました。
出力値は強化学習モデルのAction空間である -1.0 〜 1.0 の範囲に収まっており、固定入力に対して常に同じ値を返す(決定論的である)ことが確認できました。

これにより、Python環境で構築した制御モデルをSTM32マイコンへ実装する Sim-to-Realパイプライン が確立されました。

7. まとめ

本日の実装により、Python環境で構築した「脳(AIモデル)」を、C言語ベースの「神経系(マイコン)」に移植する Sim-to-Realパイプライン が完全に確立されました。

  • ONNX変換: PyTorchのTensor計算グラフを維持したままエクスポートする方法を確立。
  • コード生成: X-CUBE-AIにより、軽量かつ高速なCコードを生成。
  • 実機動作: STM32G4上で、制御周期に間に合う速度で推論できることを確認。

ここまではあくまで「計算」の世界の話です。次はいよいよ、この計算結果を「物理的な力」に変えるフェーズに入ります。

次回予告 [Day 12]

AIの実装は完了しました。次は、この推論結果(トルク指令値)を、物理的なモータの回転力に変換するための 「ベクトル制御 (FOC) の理論」 について解説します。
なぜ単なるPWMではなくFOCが必要なのか? クラーク変換・パーク変換とは何か? 数式と図解を用いて、DCブラシレスモータを静音かつ高効率に回す仕組みを紐解きます。


アドベントカレンダー参加中
STM32×AIで「3軸倒立振子」を作る25日間(ひとりアドカレ)Advent Calendar 2025

3
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
3
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?