π0シリーズで使用されるaction expertをコードレベルで理解する

![](https://cdn-ak.f.st-hatena.com/images/fotolife/S/SO0529/20251212/20251212011353.png) ABEJAアドベントカレンダー2025の24日目の記事です。 ABEJAでデータサイエンティストをしている大谷です。 最近VLAに触れ合うことが増えてきました。11/7にもMacStudioで動かすSO-ARM x π0.5 -解説から実機動作まで- というわけで今回のブログではopenpiの実装からaction_expertの部分に絞ってコードレベルで見ていくことで理解を深めていきます。 action expertとは サンプルデータでFlow Matchingを試す π0のaction_expert部分を見る 実装 action expertで使う諸々 action expert本体 サンプルデータを投入して流れを見る π0から始まったπシリーズ（勝手にシリーズと呼んでいます）は、Physical Intelligence 社が開発した VLA モデルです。 手前味噌で恐縮ですが、π0の仕組みは弊社のテックブログを参考にいただける幸いです。 https://arxiv.org/pdf/2410.24164arxiv.org その後、π0.5、π*0.6と続きます。後続モデルでもaction expertは使用されています。 https://arxiv.org/pdf/2504.16054arxiv.org action expertとは ざっくり連続したアクション生成をするために、VLMモデルに追加された部分です。ロボットはアームの回転角制御など可能な限り細かい周期で連続的に操作していきたいです。ただ出力するactionがVLMの出力速度にピン留めされてしまうと、ロボットの制御がカクついてしまいます。そこでaction expertをVLMの外に後付けし、高精度の連続アクション分布をFlow Matchingで生成することで、高い周波数での制御値を連続値として出力できるようにしました。 ちなみにFlow Matchingもわかるようでわかっていないことが多いので、ここも合わせてコードで理解していこうと思います。 Flow Matchingはざっくりいうと、適当な確率分布（ノイズ分布など）から欲しい確率分布まで、徐々に意味のある滑らかな動きに変化させる技術です。ロボットでいうとノイズから正解のアクションへ向かう中間アクションを学習する感じです。 ![](https://cdn-ak.f.st-hatena.com/images/fotolife/S/SO0529/20251211/20251211132411.png) 「Flow Matching Guide and Code」の図2 https://arxiv.org/pdf/2412.06264arxiv.org サンプルデータでFlow Matchingを試す Flow Matching Guide and CodeのP7にサンプルコードがあるので、これを使って基礎から理解します。 import torch from torch import nn, Tensor import matplotlib.pyplot as plt from sklearn.datasets import make_moons class Flow(nn.Module): def __init__(self, dim: int = 2, h: int = 64): super().__init__() self.net = nn.Sequential( nn.Linear(dim + 1, h), nn.ELU(), nn.Linear(h, h), nn.ELU(), nn.Linear(h, h), nn.ELU(), nn.Linear(h, dim) ) # x_t: (B,2), t: (B,1) def forward(self, x_t: Tensor, t: Tensor) -> Tensor: return self.net(torch.cat((x_t, t), dim=-1)) # simple midpoint ODE solver step def step(self, x_t: Tensor, t_start: float, t_end: float) -> Tensor: t_start_tensor = torch.full((x_t.shape[0], 1), t_start) t_end_tensor = torch.full((x_t.shape[0], 1), t_end) dt = t_end - t_start # midpoint method k1 = self.forward(x_t, t_start_tensor) x_mid = x_t + 0.5 * dt * k1 t_mid = torch.full((x_t.shape[0], 1), (t_start + t_end) / 2) k2 = self.forward(x_mid, t_mid) return x_t + dt * k2 モデル定義は超シンプルですね。step関数ではx_tをt_startからt_endまで1ステップ分移動させます。 dt = t_end - t_start # モデルを定義 flow = Flow() optimizer = torch.optim.Adam(flow.parameters(), lr=1e-2) loss_fn = nn.MSELoss() # 学習 for step in range(10000): # いつも三日月（ターゲット） x_1_np, _ = make_moons(256, noise=0.05) x_1 = torch.tensor(x_1_np, dtype=torch.float32) # ノイズ x_0 = torch.randn_like(x_1) # x_1とx_0を直線で結ぶ t = torch.rand(len(x_1), 1) x_t = (1 - t) * x_0 + t * x_1 # 正解の速度ベクトル dx_t = x_1 - x_0 optimizer.zero_grad() pred = flow(x_t, t) # MSEで学習 loss = loss_fn(pred, dx_t) loss.backward() optimizer.step() if step % 500 == 0: print(f"step {step}, loss={loss.item():.5f}") 上記の学習で、ノイズ状態のデータを特定の方向に輸送するベクトル場を獲得できるようになります。 ここまでくると同じような原理でロボットのアクション生成において、特定のスタート地点から次の地点までの中間経路を獲得できそうというのがわかってきます。 実際にサンプリングしたノイズから三日月データの生成過程を見ると以下のようになります。 x = torch.randn(300, 2) # t=0からt=1まで8ステップに刻む n_steps = 8 fig, axes = plt.subplots(1, n_steps + 1, figsize=(30, 4), sharex=True, sharey=True) time_steps = torch.linspace(0, 1.0, n_steps + 1) axes[0].scatter(x[:, 0].detach(), x[:, 1].detach(), s=10) axes[0].set_title(f't = {time_steps[0]:.2f}') axes[0].set_xlim(-3.0, 3.0) axes[0].set_ylim(-3.0, 3.0) for i in range(n_steps): # 現在のベクトルと次の時刻tを入力し、学習済みのベクトル場に従って次の地点までxが進む x = flow.step(x, float(time_steps[i]), float(time_steps[i + 1])) axes[i + 1].scatter(x[:, 0].detach(), x[:, 1].detach(), s=10) axes[i + 1].set_title(f't = {time_steps[i + 1]:.2f}') plt.tight_layout() plt.show() ![](https://cdn-ak.f.st-hatena.com/images/fotolife/S/SO0529/20251211/20251211212307.png) ノイズから三日月の生成過程 π0のaction_expert部分を見る openpiのレポジトリでは以下のあたりにあります。 https://github.com/Physical-Intelligence/openpi/blob/main/src/openpi/models_pytorch/pi0_pytorch.py#L100 https://github.com/Physical-Intelligence/openpi/blob/main/src/openpi/models_pytorch/pi0_pytorch.py#L376 今回はVLM側には触れず、action_expertの部分だけ見たいので、必要な実装だけ抜き出してサンプルデータで動きを見ていこうと思います。 最後に全体コードを貼りますが、長くなるので説明のため必要な部分だけを記載しています。 本家実装を見つつ、claude codeと対話しつつ、action expertの部分だけ抜き出しました。気づいたら三日月サンプルから急にややこしくなってしまってました... コード内のコメントでも何をしているのかを補足していきます。 ![](https://cdn-ak.f.st-hatena.com/images/fotolife/S/SO0529/20251224/20251224124220.png) action expertのフロー図 action expertで使う諸々 ここはaction expert特有のものでもないので割愛していきます。 @dataclass class Pi0ActionConfig: """適当な設定""" # モデル hidden_dim: int = 1024 num_layers: int = 18 mlp_dim: int = 4096 num_heads: int = 8 num_kv_heads: int = 1 head_dim: int = 256 # アクション action_dim: int = 32 # 32個の制御対象 action_horizon: int = 50 # 50ステップ先まで出力 # 時間埋め込み time_min_period: float = 4e-3 time_max_period: float = 4.0 dropout: float = 0.0 rope_theta: float = 10000.0 @property def num_kv_groups(self) -> int: return self.num_heads // self.num_kv_heads class ActionExpertAttention(nn.Module): def __init__(self, config: Pi0ActionConfig, layer_idx: int): super().__init__() （action expertと直接関係ない普通のattentionなので割愛） def forward( self, hidden_states: Tensor, attention_mask: Optional[Tensor] = None, position_ids: Optional[Tensor] = None, cos: Optional[Tensor] = None, sin: Optional[Tensor] = None, ) -> Tensor: """Forward pass Args: hidden_states: [batch, seq, hidden_dim] attention_mask: [batch, 1, seq, seq] (0 = attend, -inf = mask) position_ids: [batch, seq] cos, sin: RoPE用 [batch, seq, 1, head_dim] Returns: [batch, seq, hidden_dim] """ （省略） class ActionExpertMLP(nn.Module): def __init__(self, config: Pi0ActionConfig): super().__init__() self.gate_proj = nn.Linear(config.hidden_dim, config.mlp_dim, bias=False) self.up_proj = nn.Linear(config.hidden_dim, config.mlp_dim, bias=False) self.down_proj = nn.Linear(config.mlp_dim, config.hidden_dim, bias=False) self.act_fn = nn.GELU(approximate="tanh") def forward(self, x: Tensor) -> Tensor: return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class ActionExpertLayer(nn.Module): """Action Expert 1層分""" def __init__(self, config: Pi0ActionConfig, layer_idx: int): super().__init__() self.config = config self.input_layernorm = RMSNorm(config.hidden_dim) self.post_attention_layernorm = RMSNorm(config.hidden_dim) self.self_attn = ActionExpertAttention(config, layer_idx) self.mlp = ActionExpertMLP(config) def forward( self, hidden_states: Tensor, attention_mask: Optional[Tensor] = None, position_ids: Optional[Tensor] = None, cos: Optional[Tensor] = None, sin: Optional[Tensor] = None, ) -> Tensor: (Attention1層分の計算のみ) action expert本体 ここがVLMなしで動くaction expertの本体です。とはいえ、基本的には普通のTransformerと同じです。 ロボットの現在状態とノイズが混ざったアクションと時間を受け取って、アクションをどの方向に修正していくかを出力します。 embed_suffix class ActionExpert(nn.Module): """ VLMなしで動作する単体のaction expert 入力: - state: [B, 32] ロボット状態 - noisy_actions: [B, 50, 32] ノイズ付きアクション - timestep: [B] Flow Matchingのタイムステップ 出力: - velocity: [B, 50, 32] 予測されたベクトル場 """ def __init__(self, config: Pi0ActionConfig): super().__init__() self.config = config self.state_proj = nn.Linear(config.action_dim, config.hidden_dim) self.action_in_proj = nn.Linear(config.action_dim, config.hidden_dim) # アクション + 時間のMLP融合 self.action_time_mlp_in = nn.Linear( 2 * config.hidden_dim, config.hidden_dim ) self.action_time_mlp_out = nn.Linear(config.hidden_dim, config.hidden_dim) # RoPE self.rotary_emb = RotaryEmbedding( config.head_dim, max_seq_len=config.action_horizon + 10, theta=config.rope_theta, ) # Transformer層 self.layers = nn.ModuleList( [ActionExpertLayer(config, i) for i in range(config.num_layers)] ) # 正規化 self.norm = RMSNorm(config.hidden_dim) # 出力 self.action_out_proj = nn.Linear(config.hidden_dim, config.action_dim) def embed_suffix( self, state: Tensor, noisy_actions: Tensor, timestep: Tensor, ) -> Tuple[Tensor, Tensor]: """suffixトークンを埋め込み - state_emb: [B, 1, hidden_dim] - action_emb: [B, action_horizon, hidden_dim] + time_emb - suffix: [B, 1 + action_horizon, hidden_dim] Args: state: [batch, action_dim] noisy_actions: [batch, action_horizon, action_dim] t

π0シリーズで使用されるaction expertをコードレベルで理解する

背景や根拠まで確認しますか？

関連記事

調べる

選ぶ

サイト