• B: 배치 사이즈 (Batch Size)
• num_latest_obs (L): 2 (과거 몇 개의 관측을 볼 것인가)
• xf_n_embd (E): 256 (Transformer의 임베딩 차원)
• action_prediction_horizon (T_act): 8 (예측할 미래 행동의 길이)
• action_dim (A): 21 (행동의 차원)

Part 1: Transformer Encoder (WBVIMAPolicy.forward)

목표: 관측 데이터(obs)를 입력받아 문맥 정보가 담긴 transformer_output (Shape: [B, 6, E])을 생성.

1. 입력 데이터 (obs):
   • prop_obs: 로봇 관절 등 자체 상태. Shape: [B, L, prop_dim] = [B, 2, 21]
   • pointcloud: PCD Shape: [B, L, num_points, 6] (e.g., [B, 2, 1024, 6])
2. ObsTokenizer (관측 데이터 토큰화):
   • prop_obs는 MLP를 통과하여 Shape이 [B, 2, 21] -> [B, 2, E] = [B, 2, 256]로 변환
   • pointcloud는 PointNet을 통과하여 Shape이 [B, 2, 1024, 6] -> [B, 2, E] = [B, 2, 256]로 변환.
   • 두 결과가 순서대로 합쳐져(concat) obs_tokens가 됩니다.
     • obs_tokens Shape: [B, L * 2, E] = [B, 4, 256]
     • self.obs_tokenizer.num_tokens_per_step은 2가 됨 (proprioception 1개, pointcloud 1개).
3. tokens_in (Transformer 입력 시퀀스 구성):
   • Transformer에 넣기 위해 "관측 토큰"과 "행동을 질문할 토큰(action_readout_token)"을 번갈아 배치
   • 한 스텝(t)당 토큰 구성: [prop_token, pc_token, action_readout_token]
   • num_latest_obs가 2이므로, 전체 시퀀스는 [prop_0, pc_0, act_q_0, prop_1, pc_1, act_q_1] 형태가 됩니다.
   • n_tokens_per_step = 2(관측) + 1(질문) = 3
   • n_total_tokens = L * n_tokens_per_step = 2 * 3 = 6
   • tokens_in 최종 Shape: [B, n_total_tokens, E] = [B, 6, 256]
4. GPT (Transformer 실행):
   • tokens_in이 Transformer 모델에 입력됩니다.
   • Transformer는 Self-Attention을 통해 6개 토큰 전체의 문맥을 계산. 각 토큰은 다른 모든 토큰의 정보를 참고하여 업데이트.
   • 입력과 출력의 Shape은 동일.
   • tokens_out (== transformer_output) 최종 Shape: [B, 6, 256]

Part 2: U-Net Decoder (WBVIMAPolicy.compute_loss)

목표: transformer_output을 조건으로, gt_action의 노이즈를 예측하여 손실을 계산.

1. _get_action_readout_tokens (조건 정보 추출):
   • transformer_output([B, 6, 256])에서 행동 생성을 위해 질문했던 action_readout_token에 해당하는 결과만 추출
   • 이 토큰들은 3번째, 6번째 위치에 있음. ([prop_0, pc_0, act_q_0, prop_1, pc_1, act_q_1])
   • 추출된 action_readout_tokens Shape: [B, L, E] = [B, 2, 256]
   • 이것이 U-Net에 전달될 핵심 조건(condition) 정보입니다. 각 토큰은 "t=0 시점의 정보를 바탕으로 예측할 행동"과 "t=1 시점의 정보를 바탕으로 예측할 행동"의 문맥을 각각 담고 있습니다.