[codex] implement AdaVLA-RL adaptive compute

README.md

@@ -79,6 +79,25 @@ See [LIBERO.md](LIBERO.md) for fine-tuning/evaluating on LIBERO simulation bench
 
 See [ALOHA.md](ALOHA.md) for fine-tuning/evaluating on real-world ALOHA robot tasks.
 
+## AdaVLA-RL Adaptive Compute
+
+This branch adds AdaVLA-RL adaptive-compute training and evaluation utilities for choosing token budgets, language-model depth, and slow/fast path behavior during policy inference.
+
+Relevant entry points:
+* `vla-scripts/train_adavla_rl.py`: offline PPO-style compute-policy training with an action-error reward proxy.
+* `vla-scripts/train_adavla_rl_libero.py`: online LIBERO rollout training with task-success reward.
+* `vla-scripts/finetune.py`: OpenVLA-OFT fine-tuning with optional `--use_adaptive_compute True`.
+
+A smoke run was completed on an NVIDIA A800 80GB using `libero_spatial_no_noops` from `/root/autodl-tmp/modified_libero_rlds_sample`:
+* Training: 10 gradient steps with LoRA, L1 action head, proprio input, and adaptive compute enabled.
+* Evaluation: 10 RLDS batches with the saved LoRA adapter, action head, proprio projector, and compute policy.
+* Result summary: `mean_loss_value=1.24453125`, `mean_curr_action_l1_loss=1.1921875`, `mean_next_actions_l1_loss=1.25234375`, `mean_compute_cost=0.11611328`, `mean_compute_token_budget=32.0`, `mean_compute_layer_depth=4.0`.
+
+The remote run logs and `eval_results.json` are available in the checkpoint-free GitHub release asset:
+https://github.com/scout-123-china/openvla-oft/releases/tag/adavla-rl-remote-no-checkpoints-20260514-1430
+
+Checkpoint and model-weight files (`*.pt`, `*.pth`, `*.safetensors`) are intentionally excluded from GitHub uploads.
+
 ## Support
 
 If you run into any issues, please open a new GitHub issue. If you do not receive a response within 2 business days, please email Moo Jin Kim (moojink@cs.stanford.edu) to bring the issue to his attention.

experiments/robot/libero/run_libero_eval.py

@@ -33,6 +33,8 @@
 )
 from experiments.robot.openvla_utils import (
     get_action_head,
+    get_compute_policy,
+    get_fast_path_action_head,
     get_noisy_action_projector,
     get_processor,
     get_proprio_projector,
@@ -95,6 +97,15 @@ class GenerateConfig:
     use_film: bool = False                           # If True, uses FiLM to infuse language inputs into visual features
     num_images_in_input: int = 2                     # Number of images in the VLA input (default: 1)
     use_proprio: bool = True                         # Whether to include proprio state in input
+    use_adaptive_compute: bool = False               # If True, enables AdaVLA-RL dynamic compute at inference
+    adaptive_compute_mode: str = "greedy"            # "greedy" for eval, "sample" for exploration/debugging
+    compute_policy_checkpoint: Union[str, Path] = "" # Optional compute_policy checkpoint file or directory
+    compute_policy_hidden_dim: int = 512             # Must match training
+    compute_token_budgets: tuple = (16, 32, 64, 128) # Candidate visual token budgets
+    compute_layer_depths: tuple = (4, 8, 12)         # Candidate language-model depths
+    use_fast_path: bool = False                      # If True, load and use optional fast-path action head
+    fast_path_checkpoint: Union[str, Path] = ""      # Optional fast-path checkpoint file or directory
+    fast_path_hidden_dim: int = 512                  # Must match fast-path training
 
     center_crop: bool = True                         # Center crop? (if trained w/ random crop image aug)
     num_open_loop_steps: int = 8                     # Number of actions to execute open-loop before requerying policy
@@ -167,13 +178,21 @@ def initialize_model(cfg: GenerateConfig):
     if cfg.use_diffusion:
         noisy_action_projector = get_noisy_action_projector(cfg, model.llm_dim)
 
+    compute_policy = None
+    if cfg.use_adaptive_compute:
+        compute_policy = get_compute_policy(cfg, model.llm_dim)
+
+    fast_path_action_head = None
+    if cfg.use_fast_path:
+        fast_path_action_head = get_fast_path_action_head(cfg, model.llm_dim)
+
     # Get OpenVLA processor if needed
     processor = None
     if cfg.model_family == "openvla":
         processor = get_processor(cfg)
         check_unnorm_key(cfg, model)
 
-    return model, action_head, proprio_projector, noisy_action_projector, processor
+    return model, action_head, proprio_projector, noisy_action_projector, compute_policy, fast_path_action_head, processor
 
 
 def check_unnorm_key(cfg: GenerateConfig, model) -> None:
@@ -285,6 +304,8 @@ def run_episode(
     action_head=None,
     proprio_projector=None,
     noisy_action_projector=None,
+    compute_policy=None,
+    fast_path_action_head=None,
     initial_state=None,
     log_file=None,
 ):
@@ -336,6 +357,8 @@ def run_episode(
                     action_head=action_head,
                     proprio_projector=proprio_projector,
                     noisy_action_projector=noisy_action_projector,
+                    compute_policy=compute_policy,
+                    fast_path_action_head=fast_path_action_head,
                     use_film=cfg.use_film,
                 )
                 action_queue.extend(actions)
@@ -369,6 +392,8 @@ def run_task(
     action_head=None,
     proprio_projector=None,
     noisy_action_projector=None,
+    compute_policy=None,
+    fast_path_action_head=None,
     total_episodes=0,
     total_successes=0,
     log_file=None,
@@ -418,6 +443,8 @@ def run_task(
             action_head,
             proprio_projector,
             noisy_action_projector,
+            compute_policy,
+            fast_path_action_head,
             initial_state,
             log_file,
         )
@@ -468,7 +495,15 @@ def eval_libero(cfg: GenerateConfig) -> float:
     set_seed_everywhere(cfg.seed)
 
     # Initialize model and components
-    model, action_head, proprio_projector, noisy_action_projector, processor = initialize_model(cfg)
+    (
+        model,
+        action_head,
+        proprio_projector,
+        noisy_action_projector,
+        compute_policy,
+        fast_path_action_head,
+        processor,
+    ) = initialize_model(cfg)
 
     # Get expected image dimensions
     resize_size = get_image_resize_size(cfg)
@@ -496,6 +531,8 @@ def eval_libero(cfg: GenerateConfig) -> float:
             action_head,
             proprio_projector,
             noisy_action_projector,
+            compute_policy,
+            fast_path_action_head,
             total_episodes,
             total_successes,
             log_file,

experiments/robot/openvla_utils.py

@@ -24,6 +24,7 @@
 from prismatic.extern.hf.configuration_prismatic import OpenVLAConfig
 from prismatic.extern.hf.modeling_prismatic import OpenVLAForActionPrediction
 from prismatic.extern.hf.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor
+from prismatic.models.adaptive_compute import AdaptiveComputePolicy, FastPathActionHead, compute_policy_state_dim
 from prismatic.models.action_heads import DiffusionActionHead, L1RegressionActionHead
 from prismatic.models.film_vit_wrapper import FiLMedPrismaticVisionBackbone
 from prismatic.models.projectors import NoisyActionProjector, ProprioProjector
@@ -517,6 +518,69 @@ def get_action_head(cfg: Any, llm_dim: int) -> Union[L1RegressionActionHead, Dif
     return action_head
 
 
+def _resolve_local_or_component_checkpoint(cfg: Any, explicit_path_attr: str, file_pattern: str) -> str:
+    explicit_path = getattr(cfg, explicit_path_attr, "")
+    if explicit_path:
+        explicit_path = str(explicit_path)
+        if os.path.isfile(explicit_path):
+            return explicit_path
+        if os.path.isdir(explicit_path):
+            return find_checkpoint_file(explicit_path, file_pattern)
+        raise FileNotFoundError(f"Could not find `{explicit_path_attr}` at: {explicit_path}")
+
+    if model_is_on_hf_hub(cfg.pretrained_checkpoint):
+        raise ValueError(
+            f"AdaVLA-RL component `{file_pattern}` is not mapped for HF Hub checkpoints. "
+            f"Pass --{explicit_path_attr}=<local checkpoint file or directory>."
+        )
+    return find_checkpoint_file(cfg.pretrained_checkpoint, file_pattern)
+
+
+def get_compute_policy(cfg: Any, llm_dim: int) -> AdaptiveComputePolicy:
+    """Load AdaVLA-RL compute policy from a checkpoint."""
+
+    state_dim = compute_policy_state_dim(
+        llm_dim,
+        previous_action_dim=ACTION_DIM,
+        uncertainty_dim=1,
+        include_intermediate=True,
+    )
+    compute_policy = AdaptiveComputePolicy(
+        state_dim=state_dim,
+        hidden_dim=getattr(cfg, "compute_policy_hidden_dim", 512),
+        token_budgets=getattr(cfg, "compute_token_budgets", (16, 32, 64, 128)),
+        layer_depths=getattr(cfg, "compute_layer_depths", (4, 8, 12)),
+    ).to(DEVICE)
+
+    checkpoint_path = _resolve_local_or_component_checkpoint(cfg, "compute_policy_checkpoint", "compute_policy")
+    state_dict = load_component_state_dict(checkpoint_path)
+    compute_policy.load_state_dict(state_dict)
+    compute_policy.eval()
+    return compute_policy
+
+
+def get_fast_path_action_head(cfg: Any, llm_dim: int) -> FastPathActionHead:
+    """Load optional lightweight fast-path action head."""
+
+    state_dim = compute_policy_state_dim(
+        llm_dim,
+        previous_action_dim=ACTION_DIM,
+        uncertainty_dim=1,
+        include_intermediate=True,
+    )
+    fast_path_action_head = FastPathActionHead(
+        state_dim=state_dim,
+        hidden_dim=getattr(cfg, "fast_path_hidden_dim", 512),
+        action_dim=ACTION_DIM,
+    ).to(DEVICE)
+
+    checkpoint_path = _resolve_local_or_component_checkpoint(cfg, "fast_path_checkpoint", "fast_path")
+    state_dict = load_component_state_dict(checkpoint_path)
+    fast_path_action_head.load_state_dict(state_dict)
+    fast_path_action_head.eval()
+    return fast_path_action_head
+
+
 def resize_image_for_policy(img: np.ndarray, resize_size: Union[int, Tuple[int, int]]) -> np.ndarray:
     """
     Resize an image to match the policy's expected input size.
@@ -721,6 +785,8 @@ def get_vla_action(
     action_head: Optional[torch.nn.Module] = None,
     proprio_projector: Optional[torch.nn.Module] = None,
     noisy_action_projector: Optional[torch.nn.Module] = None,
+    compute_policy: Optional[torch.nn.Module] = None,
+    fast_path_action_head: Optional[torch.nn.Module] = None,
     use_film: bool = False,
 ) -> List[np.ndarray]:
     """
@@ -735,6 +801,8 @@ def get_vla_action(
         action_head: Optional action head for continuous actions
         proprio_projector: Optional proprioception projector
         noisy_action_projector: Optional noisy action projector for diffusion
+        compute_policy: Optional AdaVLA-RL compute policy
+        fast_path_action_head: Optional lightweight fast-path action head
         use_film: Whether to use FiLM
 
     Returns:
@@ -780,7 +848,15 @@ def get_vla_action(
         # Generate action
         if action_head is None:
             # Standard VLA output (single-image inputs, discrete actions)
-            action, _ = vla.predict_action(**inputs, unnorm_key=cfg.unnorm_key, do_sample=False)
+            action, _ = vla.predict_action(
+                **inputs,
+                unnorm_key=cfg.unnorm_key,
+                do_sample=False,
+                compute_policy=compute_policy,
+                fast_path_action_head=fast_path_action_head,
+                use_adaptive_compute=getattr(cfg, "use_adaptive_compute", False),
+                adaptive_compute_mode=getattr(cfg, "adaptive_compute_mode", "greedy"),
+            )
         else:
             # Custom action head for continuous actions
             action, _ = vla.predict_action(
@@ -792,6 +868,10 @@ def get_vla_action(
                 noisy_action_projector=noisy_action_projector,
                 action_head=action_head,
                 use_film=use_film,
+                compute_policy=compute_policy,
+                fast_path_action_head=fast_path_action_head,
+                use_adaptive_compute=getattr(cfg, "use_adaptive_compute", False),
+                adaptive_compute_mode=getattr(cfg, "adaptive_compute_mode", "greedy"),
             )
 
     # Return action chunk as list of actions

experiments/robot/robot_utils.py

@@ -105,6 +105,8 @@ def get_action(
     action_head: Optional[torch.nn.Module] = None,
     proprio_projector: Optional[torch.nn.Module] = None,
     noisy_action_projector: Optional[torch.nn.Module] = None,
+    compute_policy: Optional[torch.nn.Module] = None,
+    fast_path_action_head: Optional[torch.nn.Module] = None,
     use_film: bool = False,
 ) -> Union[List[np.ndarray], np.ndarray]:
     """
@@ -119,6 +121,8 @@ def get_action(
         action_head: Optional action head for continuous actions
         proprio_projector: Optional proprioception projector
         noisy_action_projector: Optional noisy action projector for diffusion
+        compute_policy: Optional AdaVLA-RL compute policy
+        fast_path_action_head: Optional lightweight fast-path action head
         use_film: Whether to use FiLM
 
     Returns:
@@ -138,6 +142,8 @@ def get_action(
                 action_head=action_head,
                 proprio_projector=proprio_projector,
                 noisy_action_projector=noisy_action_projector,
+                compute_policy=compute_policy,
+                fast_path_action_head=fast_path_action_head,
                 use_film=use_film,
             )
         else: