code/src/models/rt2_like_baseline.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoTokenizer,
    SiglipVisionModel, SiglipImageProcessor, LlamaForCausalLM, 
    LogitsProcessor, LogitsProcessorList
)
from PIL import Image

from .encoder import ActionTransformerProjector


class LlamaProcessorWrapper:
    def __init__(self, tokenizer, image_processor):
        self.tokenizer = tokenizer
        self.image_processor = image_processor

class RT2LikeBaseline(nn.Module):
    def __init__(
        self, 
        lmm_path="meta-llama/Llama-3.2-1B-Instruct",
        vision_encoder_path="google/siglip2-base-patch16-256",
        action_dim=7,
        num_actions=1,
        num_history=0,
        use_proprio_input_vlm=True,
        use_transformer_projector=True,
        projector_depth=2,
        projector_num_heads=4,
        backbone_mode="finetune", # Options: "frozen", "finetune"
        gradient_checkpointing=True,
        num_bins=256,
        attn_implementation="flash_attention_2",
    ):
        super().__init__()
        
        print(f"Initializing RT-2 Baseline with attn_implementation: {attn_implementation}")

        self.lmm = LlamaForCausalLM.from_pretrained(
            lmm_path, dtype=torch.bfloat16, attn_implementation=attn_implementation
        )
        self.vision_encoder = SiglipVisionModel.from_pretrained(
            vision_encoder_path, dtype=torch.bfloat16, attn_implementation=attn_implementation
        )
        
        tokenizer = AutoTokenizer.from_pretrained(lmm_path)
        if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
        image_processor = SiglipImageProcessor.from_pretrained(vision_encoder_path)
        self.processor = LlamaProcessorWrapper(tokenizer, image_processor)

        self.hidden_size = self.lmm.config.hidden_size
        self.vision_projector = nn.Sequential(
            nn.Linear(self.vision_encoder.config.hidden_size, self.hidden_size),
            nn.SiLU(), 
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.SiLU(), 
            nn.Linear(self.hidden_size, self.hidden_size)
        )

        if backbone_mode == "frozen":
            self.lmm.requires_grad_(False)
            self.vision_encoder.requires_grad_(False)
        elif backbone_mode == "finetune":
            self.lmm.requires_grad_(True)
            self.vision_encoder.requires_grad_(True)
        else:
            raise ValueError(f"Unknown backbone_mode: {backbone_mode}")

        if gradient_checkpointing:
            if hasattr(self.lmm, "gradient_checkpointing_enable"):
                self.lmm.gradient_checkpointing_enable()
            if hasattr(self.lmm, "enable_input_require_grads"):
                self.lmm.enable_input_require_grads()
            if hasattr(self.lmm.config, "use_cache"):
                self.lmm.config.use_cache = False
            if hasattr(self.vision_encoder, "gradient_checkpointing_enable"):
                self.vision_encoder.gradient_checkpointing_enable()

        self.action_dim = action_dim
        self.num_actions = num_actions
        self.num_history = num_history
        self.num_bins = num_bins
        self.use_proprio_input_vlm = use_proprio_input_vlm

        if self.use_proprio_input_vlm:
            if use_transformer_projector:
                self.action_projector = ActionTransformerProjector(
                    action_dim=action_dim,
                    hidden_size=self.hidden_size,
                    depth=projector_depth,
                    num_heads=projector_num_heads
                )
            else:
                self.action_projector = nn.Linear(action_dim, self.hidden_size)
        else:
            self.action_projector = None

        vocab_limit = 128000
        self.action_token_start_idx = vocab_limit - num_bins
        print(f"RT-2 Mode: Using vocabulary indices [{self.action_token_start_idx}, {vocab_limit}) for {num_bins} action bins.")

    def forward(self, input_ids=None, attention_mask=None, actions=None, proprioception=None, proprio_attention_mask=None, pixel_values=None, **kwargs):
        vision_outputs = self.vision_encoder(pixel_values, output_hidden_states=True)
        image_feats = vision_outputs.last_hidden_state
        image_embeds = self.vision_projector(image_feats)

        B = input_ids.shape[0]

        if image_embeds.shape[0] != B:
            num_views = image_embeds.shape[0] // B
            image_embeds = image_embeds.view(B, num_views, -1, image_embeds.shape[-1])
            image_embeds = image_embeds.flatten(1, 2)

        text_embeds = self.lmm.model.embed_tokens(input_ids)

        proprio_embeds = None
        if self.use_proprio_input_vlm and proprioception is not None:
            proprio_embeds = self.action_projector(proprioception.to(device=text_embeds.device, dtype=text_embeds.dtype))

        actions_flat = actions.view(B, -1)
        
        gt_actions_clamped = torch.clamp(actions_flat, -1, 1)
        gt_actions_idx = ((gt_actions_clamped + 1) / 2 * (self.num_bins - 1)).round().long()
        
        action_input_ids = gt_actions_idx + self.action_token_start_idx
        action_embeds = self.lmm.model.embed_tokens(action_input_ids)

        embeds_list = [image_embeds]
        
        image_mask = torch.ones(B, image_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
        mask_list = [image_mask]
        
        labels_list = [torch.full((B, image_embeds.shape[1]), -100, dtype=torch.long, device=attention_mask.device)]

        if proprio_embeds is not None:
            embeds_list.append(proprio_embeds)
            
            p_len = proprio_embeds.shape[1]
            if proprio_attention_mask is not None:
                mask_list.append(proprio_attention_mask.to(attention_mask.device))
            else:
                p_mask = torch.ones(B, p_len, device=attention_mask.device, dtype=attention_mask.dtype)
                mask_list.append(p_mask)
            
            labels_list.append(torch.full((B, p_len), -100, dtype=torch.long, device=attention_mask.device))
        
        embeds_list.append(text_embeds)
        mask_list.append(attention_mask)
        labels_list.append(torch.full((B, text_embeds.shape[1]), -100, dtype=torch.long, device=attention_mask.device))

        embeds_list.append(action_embeds)
        action_mask = torch.ones(B, action_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)
        mask_list.append(action_mask)
        labels_list.append(action_input_ids)

        inputs_embeds = torch.cat(embeds_list, dim=1)
        combined_attention_mask = torch.cat(mask_list, dim=1)
        combined_labels = torch.cat(labels_list, dim=1)

        outputs = self.lmm(
            inputs_embeds=inputs_embeds,
            attention_mask=combined_attention_mask,
            labels=combined_labels,
        )
        
        return outputs.loss

    @torch.no_grad()
    def predict_action(self, input_ids, attention_mask, proprioception=None, proprio_attention_mask=None, pixel_values=None, **kwargs):
        B = input_ids.shape[0]

        vision_outputs = self.vision_encoder(pixel_values, output_hidden_states=True)
        image_embeds = self.vision_projector(vision_outputs.last_hidden_state)

        if image_embeds.shape[0] != B:
            num_views = image_embeds.shape[0] // B
            image_embeds = image_embeds.view(B, num_views, -1, image_embeds.shape[-1])
            image_embeds = image_embeds.flatten(1, 2)

        text_embeds = self.lmm.model.embed_tokens(input_ids)

        embeds_list = [image_embeds]
        mask_list = [torch.ones(B, image_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype)]
        
        if self.use_proprio_input_vlm and proprioception is not None:
            proprio_embeds = self.action_projector(proprioception)
            embeds_list.append(proprio_embeds)
            if proprio_attention_mask is not None:
                mask_list.append(proprio_attention_mask.to(attention_mask.device))
            else:
                mask_list.append(torch.ones(B, proprio_embeds.shape[1], device=attention_mask.device, dtype=attention_mask.dtype))
        
        embeds_list.append(text_embeds)
        mask_list.append(attention_mask)

        inputs_embeds = torch.cat(embeds_list, dim=1)
        attention_mask = torch.cat(mask_list, dim=1)

        class RT2ActionLogitsProcessor(LogitsProcessor):
            def __init__(self, start_idx, end_idx):
                self.start_idx = start_idx
                self.end_idx = end_idx
            
            def __call__(self, input_ids, scores):
                scores[:, :self.start_idx] = float('-inf')
                scores[:, self.end_idx:] = float('-inf')
                return scores

        logits_processor = LogitsProcessorList([
            RT2ActionLogitsProcessor(self.action_token_start_idx, self.action_token_start_idx + self.num_bins)
        ])
        
        total_action_tokens = self.num_actions * self.action_dim

        pad_token_id = getattr(self.lmm.config, "pad_token_id", None)
        pad_token_id = pad_token_id if pad_token_id is not None else 0

        dummy_input_ids = torch.full(
            (B, inputs_embeds.shape[1]), 
            pad_token_id,
            dtype=torch.long, 
            device=inputs_embeds.device
        )

        outputs = self.lmm.generate(
            input_ids=dummy_input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            max_new_tokens=total_action_tokens,
            logits_processor=logits_processor,
            do_sample=False,
            use_cache=True,
            eos_token_id=[],
            pad_token_id=pad_token_id
        )

        generated_ids = outputs[:, inputs_embeds.shape[1]:]

        bin_indices = generated_ids - self.action_token_start_idx
        
        action_flat = (bin_indices.float() / (self.num_bins - 1)) * 2 - 1
        action = action_flat.view(B, self.num_actions, self.action_dim).to(dtype=self.lmm.dtype)
        
        return action


if __name__ == "__main__":
    print("Testing RT-2 Like Baseline...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16
    
    # Initialize Model
    model = RT2LikeBaseline(
        lmm_path="meta-llama/Llama-3.2-1B-Instruct",
        vision_encoder_path="google/siglip2-base-patch16-256",
        action_dim=7, num_actions=8, num_history=8,
        backbone_mode="finetune", gradient_checkpointing=False
    ).to(device, dtype)

    B = 2
    # Dummy Data
    img = Image.new('RGB', (256, 256), color='red')
    imgs = [img] * B
    
    # Process images
    pixel_values = torch.stack([
        model.processor.image_processor(i, return_tensors="pt")["pixel_values"].squeeze(0) for i in imgs
    ]).to(device, dtype)
    
    # Process text
    texts = ["Pick up the red block."] * B
    text_inputs = model.processor.tokenizer(texts, padding=True, return_tensors="pt")
    input_ids = text_inputs["input_ids"].to(device)
    attention_mask = text_inputs["attention_mask"].to(device)
    
    # Dummy tensors
    act_gt = torch.randn(B, 8, 7, device=device, dtype=dtype)
    proprio = torch.randn(B, 8, 7, device=device, dtype=dtype)

    # Test forward (training)
    loss = model(input_ids=input_ids, attention_mask=attention_mask, actions=act_gt,
                 proprioception=proprio, pixel_values=pixel_values)
    print(f"Training Loss: {loss.item():.4f}")
    
    # Test predict_action (inference)
    pred = model.predict_action(input_ids=input_ids, attention_mask=attention_mask,
                                proprioception=proprio, pixel_values=pixel_values)
    print(f"Predicted Action Shape: {pred.shape}")
    
    print("\nTest Passed!")