技術概要：MuJoCoUni

問題定義

ロボット学習システムは、物理シミュレータをトレーニングループに直接統合する傾向が強まっており、高頻度でのバッチ処理による物理ステップ、センサー読み取り、環境リセットの実行を必要としています。アップストリームの MuJoCo は成熟したアセットとソルバを提供していますが、その標準的な Python インターフェースは、高頻度におけるオブジェクトのライフサイクル管理、インターフェース呼び出し、出力の整形に関してオーバーヘッドを発生させます。

既存のソリューションにはトレードオフが存在します：

• アップストリームの mujoco.rollout： オープンループの軌道生成（計画、システム同定）には効率的ですが、オンライン強化学習（RL）に必要な状態を保持する環境の永続性、スパースなリセット機能、および環境ごとのランダム化ライフサイクルが欠如しています。
• GPU 常駐シミュレータ（例：Brax、MJX、Isaac Gym）： 大規模な並列性を提供しますが、モデルやデータ配置が特定のアクセラレータの制約に適合することを要求し、完全な機能カバレッジ、複雑な接触処理、または既存の CPU 最適化 MuJoCo アセットの再利用能力を犠牲にする可能性があります。
• 標準的な CPU ベクトル化ランタイム： 特定の MuJoCo ソルバのセマンティクスとの深い統合が欠けているか、モデルの永続性とドメインランダム化を効率的に維持するために多大なエンジニアリングを必要とすることが多いです。

アップストリームの MuJoCo（モデル、センサー、接触、制約）の完全なセマンティック忠実性を維持しつつ、オンライン RL とバッチ物理評価のための高スループットな並列実行と状態を保持する環境管理を提供する、CPU ベースのバッチランタイムが必要です。

手法

MuJoCoUni は、MuJoCo の軽量なダウンストリーム配布版であり、エンジニアリング上の追加を厳密にPython バインディング層内に集中させ、コア物理カーネル、ソルバ、積分器、接触モデルは変更しません。

コアアーキテクチャ：BatchEnvPool

中心的なコンポーネントは BatchEnvPool であり、これは C++/pybind11 によるエグゼキュータで、以下を管理します：

1. 永続リソース： 各環境ごとに 1 つの mjModel コピー（計 $N$ 個）と、各スレッドごとに 1 つの mjData ワーカー（計 $W$ 個）を所有します。
2. スレッドプール： 内部のスレッドプールが、並列実行のためにワーカースレッドへ環境インデックスをディスパッチします。
3. 状態保持実行： 呼び出しごとに状態を持たない rollout と異なり、BatchEnvPool は呼び出し間での環境の状態を維持し、スパースなリセットと永続的なモデルバリアントを可能にします。

主要プリミティブ

API は、オンラインロボットトレーニングと評価向けに設計されたいくつかのプリミティブを公開します：

• step： すべての環境で $n$ 回の積分ステップ（mj_step）を実行します。完全な軌道の返却というオーバーヘッドを避けるため、最終状態（およびオプションで最終センサーデータ）のみを返します。
• forward： 力学を進めることなく現在の状態に基づいてセンサーを更新するために mj_forward を実行し、観測の構築に有用です。
• reset： スパースなリセットを実装し、終了した環境のサブセット（env_ids）のみに対して作用します。リセットライフサイクルによるドメインランダム化をサポートし、影響を受ける環境に対してのみフィールドパッチ（質量、慣性、摩擦など）を適用し、mj_setConst を呼び出します。
• compute_site_jacobians： 特定のサイトに対する並進および回転ヤコビアンをバッチ計算します。最小限の運動学プレフィックス（mj_kinematics、mj_comPos）を実行した後、mj_jacSite を実行します。
• sample_hfield_height： 各種フレームアライメント（ワールド、ボディ、ヨー）をサポートし、バッチで MuJoCo の高さフィールド幾何学を線形補間してサンプリングします。

設計境界

実装は MuJoCo ソースツリーの変更を避けます。環境のインスタンス化には mj_copyModel を、ランダム化には mj_setConst に依存します。設計は 2 つのモードをサポートします：

1. 単一共有モデル： すべての環境がモデルコピーを共有します（メモリ効率が良いが、ランダム化の柔軟性は低い）。
2. モデルバリアント： 各環境が独自のモデルコピーを所有し、構築時にプリコンパイルされた幾何学レベルのランダム化（リンク長、メッシュスケールなど）を可能にします。

主要な貢献

1. 状態保持バッチランタイム： オープンループ軌道生成とオンライン・状態保持 RL 環境の要件の間のギャップを埋める BatchEnvPool を導入します。
2. 効率的なスパースリセットとランダム化： 終了した環境のみをリセットし、フィールドパッチとモデルバリアントの両方によるドメインランダム化を最小限のオーバーヘッドで適用するメカニズムを提供します。これはリセットされる環境の数に対して線形にスケーリングします。
3. バッチ物理クエリ： 力学を進めることなく、センサーのフォワードパス、サイトヤコビアン、高さフィールドサンプリングの高性能なバッチ実装を公開します。
4. セマンティクスの維持： 物理挙動がアップストリームの CPU MuJoCo と同一であることを保証し、既存の XML/MJB アセット、デバッグツール、センサー定義との互換性を維持します。

結果

著者は、16 スレッドの Intel i9-14900HX CPU を使用して、4 つのロボットモデル（Unitree Go1、Wonik Allegro、Franka Panda、CMU Humanoid）で MuJoCoUni を検証しました。

• スループット： スループットは約 256〜512 環境で飽和します。飽和時、システムは以下を達成します：
  • Wonik Allegro： 約 180 万ステップ/秒。
  • Unitree Go1： 約 120 万ステップ/秒。
  • Franka Panda： 約 41 万ステップ/秒。
  • CMU Humanoid： 約 29 万ステップ/秒。
• リセット性能： C++ の BatchEnvPool パスは、Python ループよりも著しく高速です。4096 環境の完全なリセットにおいて、C++ では3.5 ms、Python の for ループでは53 msを要し（約 15 倍の高速化）。
• ヤコビアン性能： 4096 環境に対するサイトヤコビアンの計算は、C++ では0.53 ms、Python では11.9 msを要し（約 22 倍の高速化）。
• 高さフィールドサンプリング： 4096 環境に対する高さフィールドのサンプリングは、C++ では0.52 ms、Python では290 msを要し（約 555 倍の高速化）。
• モデルバリアント： 環境ごとの固有モデルコピーを使用すること（モデルバリアントモード）によるオーバーヘッドは、飽和時には単一共有モデルと比較して無視でき、スループットは実質的に同一になります。

意義と範囲

本論文は、MuJoCoUni を GPU 常駐シミュレータやアップストリームの rollout の代替ではなく、補完的なエンジニアリングソリューションとして位置づけています。

• GPU への補完： 生きたアクセラレータ常駐よりも機能カバレッジ（完全な MuJoCo セマンティクス、複雑な接触、デバッグ）とクロスプラットフォームの移植性が優先されるワークロードを対象としています。これにより、システムは CPU MuJoCo ソルバの安定性を維持しつつ、現代のマルチコア CPU を活用して高スループットなデータ生成を行うことができます。
• rollout への補完： 永続的な環境、スパースなリセット、現在状態のクエリを必要とするオンラインロボット RLと状態保持評価ループの特定のニーズに対応します。これらは、完全な軌道生成向けに設計された rollout には最適化されていないタスクです。
• システム統合： 物理実行を上位のトレーニングロジック、ログ記録、スケジューリングから分離するシミュレーションバックエンドとして機能します。ドメインランダム化、地形を認識した歩行、バッチ最適化（例：進化戦略）を必要とするシミュレーションから実世界へのトレーニングに特に適しています。

著者は、この貢献がエンジニアリング指向であり、アップストリームの MuJoCo エコシステムを維持する CPU ベースのバッチ実行のための具体的で再現可能な道筋を提供することを強調しています。