アーカイブ / Robotics
TD-MPC2の報酬予測を317Mから1Mのモデルへ移す多課題学習
Reward Distillation for TD-MPC2
Reward Distillation for TD-MPC2の原論文から掲載した図です。細部はクリックして拡大できます。
原論文の図の説明を読む
Figure 1. Our distillation approach consists of two main loss function components: the original loss from TD-MPC2 (in red), calculated as a linear combination of consistency, reward, and value losses ( α \alpha denotes a scaling coefficient for each respective loss); and the distillation loss (in teal) that calculates MSE between student’s (green) and teacher’s (blue) rewards produced from inferring the same state-action pair. The total loss is a linear combination of both losses, with the distillation loss scaled by the d_coef . In our scenario, the student is trainable while the teacher’s weights are frozen. The figure depicts a diagram of our proposed distillation pipeline. A pretrained teacher and a trainable student are used to infer rewards for reward distillation loss calculation. The original loss from TD-MPC2 is preserved, and the total loss is a linear combination of both losses, with the distillation loss scaled by the \textit{d\_coef}.
概要
- 317MパラメータのTD-MPC2を教師に、報酬予測の差を使って1Mパラメータの生徒を追加学習する。
- MT30の掲載表では、バッチサイズ256の20万ステップで17.85対14.04、100万ステップで28.12対27.36。ただしバッチサイズ1024では比較対象が上回る条件もある。
- 著者らはFP16量子化による容量50％削減を報告する。概要のスコア28.45と掲載表の28.12の条件差、および実機での有効性は確認できない。
補足：編集した模式図で仕組みを確認する
VISUAL EXPLAINER
図でつかむ、Reward Distillation for TD-MPC2の仕組み
- 01MT30の状態・行動
- 02教師と生徒が報酬予測
- 03予測差を学習に追加
- 04生徒をFP16量子化
原論文の説明をもとにした模式図です。処理の細部は省略しています。
研究の背景
この研究は、複数の課題を扱うTD-MPC2のモデルを小さくする。TD-MPC2は、環境の変化を予測するモデルを行動学習に使う、モデルベース強化学習の手法である。大きなモデルは、計算資源が限られた場所への配置が難しい。
著者らは317Mパラメータのモデルから、1Mパラメータのモデルへ知識を移す。パラメータは学習で調整される数値で、その数はモデルの保存や計算の負担に関わる。
手法
MT30で学習済みの317MパラメータのTD-MPC2を教師として固定し、1Mパラメータのチェックポイントを生徒として追加学習する。両者に同じ状態・行動を入力して報酬予測を得る。報酬は行動の結果を評価する学習上の値であり、二つの予測の差を平均二乗誤差で測る。平均二乗誤差は、差を二乗して平均した値である。
生徒の更新には、元のTD-MPC2の整合性・報酬・価値の損失も使う。損失は予測と学習目標のずれを示す値で、追加した予測差には係数を掛けて重みを調整する。例えば同じ状態・行動への報酬予測が違えば、生徒を教師の予測に近づける。学習後、生徒にFP16量子化を適用する。
新規性
著者らは、TD-MPC2の従来の学習損失を保ち、教師と生徒の報酬予測の差を新たな損失として加えた。教師の内部表現全体を生徒に再現させる設計ではない。
学習後にはFP16量子化も適用する。FP16は数値を16ビットで表す方式で、著者らはモデル容量を50％減らして性能を維持したと述べる。ただし、量子化前後の対応するスコア表は確認できた本文の範囲にはない。
従来手法との違い
比較対象は、同じ1MパラメータのTD-MPC2を教師からの蒸留なしで学習する条件である。バッチサイズ256では、20万ステップで蒸留した生徒が17.85、比較対象が14.04だった。バッチサイズは一度の更新に使うデータ数を指す。
100万ステップでは28.12対27.36と差が縮まる。一方、バッチサイズ1024では20万ステップと33万7000ステップのどちらも比較対象の方が高い。蒸留が常に優れるという結果ではない。
実験結果
著者らはMT30の全69万エピソード、3億4569万遷移を用い、正規化スコアで評価した。エピソードは一連の試行、遷移は状態から次の状態への移り変わりを指す。実験には単一のRTX 3060、12GB VRAMを使用したと記す。
掲載表では、317Mの教師を使い、バッチサイズ256で20万ステップ学習した生徒が17.85、蒸留なしの比較対象が14.04だった。100万ステップでは28.12対27.36。概要にある28.45、および比較先の18.93との条件の対応は、確認できた本文の範囲では分からない。
応用の可能性
編集上の応用案 · 論文が実証した用途とは区別しています。
編集上の応用案：複数の物体操作を行うロボット向けに、制御モデルの小型化を試す。入力として、各操作の状態・行動・次の状態・報酬を記録した学習データ、課題ラベル、学習済みの大きな教師モデルを用意する。
実機のセンサー値を学習時の状態表現に対応付け、教師と生徒に同じ状態・行動を与えて生徒を調整する。出力は課題ごとに行動を選ぶ小型モデルとする。実機での精度や速度の改善は、この論文の結果からは確定できない。
応用先の候補からNauhに相談できます。
導入時の検証
導入時の検証案 · 対象データでの再評価が必要です。
導入時の検証案：対象課題の状態・行動・次の状態・報酬を記録し、各記録に課題ラベルを付ける。実機で使うなら必要な状態を取得するセンサーを定め、その値を学習データの状態表現に合わせる。学習済みの教師と小型の生徒を用意し、対象課題の記録で生徒を調整する。
比較基準は、同じ容量の生徒を蒸留なしで学習したモデルとする。課題別成績と平均の正規化スコアを同じ条件で比べ、量子化前後の容量、推論時間、メモリ使用量も測る。実機での確認は記録データ上の評価とは別に行う。
限界と課題
著者らは、実機での導入検証が必要なこと、大きな教師モデルに依存すること、評価対象がMT30に限られることを限界として挙げる。次の状態の潜在表現も蒸留する試みでは、教師と生徒の表現の次元差があり、成績が低かった。
編集上の検証事項は、FP16量子化前後の同一条件での成績、推論時間、メモリ使用量である。概要の28.45と表の100万ステップの28.12は一致せず、条件差は確認できた本文の範囲では特定できない。
出典
元のタイトル・要旨を確認する（英語）
Knowledge Transfer in Model-Based Reinforcement Learning Agents for Efficient Multi-Task Learning
We propose an efficient knowledge transfer approach for model-based reinforcement learning, addressing the challenge of deploying large world models in resource-constrained environments. Our method distills a high-capacity multi-task agent (317M parameters) into a compact 1M parameter model, achieving state-of-the-art performance on the MT30 benchmark with a normalized score of 28.45, a substantial improvement over the original 1M parameter model's score of 18.93. This demonstrates the ability of our distillation technique to consolidate complex multi-task knowledge effectively. Additionally, we apply FP16 post-training quantization, reducing the model size by 50% while maintaining performance. Our work bridges the gap between the power of large models and practical deployment constraints, offering a scalable solution for efficient and accessible multi-task reinforcement learning in robotics and other resource-limited domains.
論文本体の取得範囲に基づく解説。AIが作成した未校閲の記事です。性能の数値は著者の評価条件に依存します。応用例と検証計画は編集上の提案です。解説更新：2026-10-01