Training › Reinforcement Learning¶

def soft_update_target(source: nn.Module, target: nn.Module, tau: float = 0.005):
    """Polyak averaging: target = (1 - tau) * target + tau * source."""
    with torch.no_grad():
        for sp, tp in zip(source.parameters(), target.parameters()):
            tp.data.mul_(1.0 - tau).add_(sp.data, alpha=tau)

def pretrain_encoder_step(
    frozen_vla,
    enc_dec: ActionTokenEncoderDecoder,
    batch_images: list,
    instructions: list,
    device: str = "cuda",
):
    """
    One pretraining step: VLA forward → encoder → decoder → reconstruction loss.

    Returns:
        recon_loss: scalar tensor (with grad on enc_dec params only)
    """
    with torch.no_grad():
        action_queries = frozen_vla.get_action_queries(
            batch_images=batch_images,
            instructions=instructions,
        )  # (B, chunk_len, H) on device

    _, recon_loss = enc_dec(action_queries)
    return recon_loss

def collect_observations_fast(
    suite_name: str,
    task_id: int,
    n_observations: int,
    steps_per_env: int = 20,
    num_envs: int = 4,
    n_initial_states: int = 50,
    libero_python: str = None,
    seed: int = 42,
) -> list:
    """
    Lightweight observation collection for encoder pretraining.

    Instead of running full episodes (300+ steps each, 99% CPU idle),
    just reset envs to diverse initial states and take a few random steps.
    Returns list of (images, instruction) tuples ready for VLA forward.

    Much faster than collect_group: no VLA forward during collection,
    no action caching, just random exploration for observation diversity.
    """
    import threading
    from concurrent.futures import ThreadPoolExecutor, as_completed

    observations = []

    # Shared progress counter across threads
    obs_per_reset = 1 + steps_per_env
    n_resets = max(1, (n_observations + obs_per_reset - 1) // obs_per_reset)
    resets_per_env = (n_resets + num_envs - 1) // num_envs
    progress_lock = threading.Lock()
    progress = {"resets_done": 0, "obs_count": 0}
    log_interval = max(1, n_resets // 20)  # log every ~5%

    logger.info(f"  [collect_obs] Plan: {n_resets} env resets × {steps_per_env} steps/reset "
                f"= ~{n_resets * obs_per_reset} obs, using {num_envs} envs")

    def _collect_from_env(env_idx, states_to_visit):
        local_rng = np.random.RandomState(seed + env_idx * 10000)
        local_obs = []
        env = LiberoEnv(libero_python=libero_python)
        try:
            for s_idx in states_to_visit:
                obs = env.reset(
                    suite_name=suite_name,
                    task_id=task_id,
                    initial_state_idx=s_idx % n_initial_states,
                    seed=seed + env_idx * 1000 + s_idx,
                )
                task_desc = env.task_description
                local_obs.append((
                    [obs["primary_image"].copy(), obs["wrist_image"].copy()],
                    task_desc,
                ))
                for _ in range(steps_per_env):
                    random_action = local_rng.uniform(-1, 1, size=7).astype(np.float32)
                    random_action[6] = local_rng.choice([-1.0, 1.0])
                    obs, _, done = env.step(random_action)
                    local_obs.append((
                        [obs["primary_image"].copy(), obs["wrist_image"].copy()],
                        task_desc,
                    ))
                    if done:
                        break
                # Update shared progress
                with progress_lock:
                    progress["resets_done"] += 1
                    rd = progress["resets_done"]
                    if rd % log_interval == 0 or rd == n_resets:
                        est_obs = rd * obs_per_reset
                        logger.info(f"  [collect_obs] reset {rd}/{n_resets} "
                                    f"({rd * 100 // n_resets}%), ~{est_obs} obs collected")
        finally:
            env.close()
        return local_obs

    with ThreadPoolExecutor(max_workers=num_envs) as pool:
        futures = {}
        for e in range(num_envs):
            start_state = e * resets_per_env
            end_state = min(start_state + resets_per_env, n_resets)
            if start_state >= n_resets:
                break
            states = list(range(start_state, end_state))
            futures[pool.submit(_collect_from_env, e, states)] = e

        for fut in as_completed(futures):
            local_obs = fut.result()
            observations.extend(local_obs)

    # Trim to requested size
    if len(observations) > n_observations:
        np.random.shuffle(observations)
        observations = observations[:n_observations]

    logger.info(f"  [collect_obs] Collected {len(observations)} observations")
    return observations

def extract_action_queries_from_obs(
    frozen_vla,
    observations: list,
    batch_size: int = 16,
    device: str = "cuda",
) -> torch.Tensor:
    """
    Batch-extract action_queries from observations via frozen VLA.

    Args:
        observations: list of (images, instruction) where images=[primary, wrist]

    Returns:
        all_queries: (N, chunk_len, H) tensor on device
    """
    N = len(observations)
    n_batches = (N + batch_size - 1) // batch_size
    all_queries = []
    for b_idx, start in enumerate(range(0, N, batch_size)):
        end = min(start + batch_size, N)
        batch_imgs = [observations[i][0] for i in range(start, end)]
        batch_instr = [observations[i][1] for i in range(start, end)]
        with torch.no_grad():
            aq = frozen_vla.get_action_queries(
                batch_images=batch_imgs,
                instructions=batch_instr,
            )  # (B, chunk_len, H)
        all_queries.append(aq)
        if (b_idx + 1) % max(1, n_batches // 10) == 0 or b_idx == n_batches - 1:
            logger.info(f"  [extract] batch {b_idx + 1}/{n_batches} "
                        f"({end}/{N} samples)")

    return torch.cat(all_queries, dim=0)  # (N, chunk_len, H)

def extract_action_queries_dataset(
    frozen_vla,
    episodes: list,
    batch_size: int = 16,
    device: str = "cuda",
) -> torch.Tensor:
    """
    Batch-extract all action_queries from rollout episodes via frozen VLA.
    (Legacy: used when full episodes are already collected.)
    """
    obs_list = []
    for ep in episodes:
        for step in ep.step_records:
            obs_list.append(
                ([step.primary_image, step.wrist_image], step.instruction)
            )
    return extract_action_queries_from_obs(frozen_vla, obs_list, batch_size, device)

@torch.no_grad()
def action_token_collect_group(
    frozen_vla,
    encoder: ActionTokenEncoderDecoder,
    actor: ActionTokenActor,
    critic: ActionTokenCritic,
    suite_name: str,
    task_id: int,
    n_initial_states: int,
    action_norm_stats: dict,
    max_steps: int,
    chunk_len: int,
    G: int = 8,
    libero_python: Optional[str] = None,
    seed: int = 42,
    num_steps_wait: int = 10,
    device: str = "cuda",
    video_dir: Optional[str] = None,
    num_envs: int = 4,
    group_idx: int = 0,
    batch_server: Optional[BatchInferenceServer] = None,
    store_images: bool = False,
    group_size: int = 1,
    reward_coef: float = 1.0,
) -> List[ActionTokenEpisode]:
    """
    Collect G episodes using RLActionToken policy.

    Uses BatchInferenceServer for GPU inference: all num_envs env threads submit
    requests concurrently; a single background thread batches them into one GPU
    forward pass. This maximizes GPU utilization vs. the old model_lock (batch=1).

    Args:
        batch_server: shared server for this GPU (created by caller for cross-task
                      batching). If None, creates a local server for this call only.
        group_size: number of trajectories per initial state. G episodes are
                    split into G//group_size unique states, each repeated
                    group_size times. Default 1 = legacy behavior (no repeat).
    """
    from concurrent.futures import ThreadPoolExecutor, as_completed

    frozen_vla.eval()
    encoder.eval()
    actor.eval()
    critic.eval()

    # Use caller-provided server (shared across tasks on same GPU) or create locally
    _own_server = batch_server is None
    if _own_server:
        batch_server = BatchInferenceServer(
            frozen_vla=frozen_vla,
            encoder=encoder,
            actor=actor,
            critic=critic,
            device=device,
        ).start()

    n_workers = min(G, num_envs)
    # Each episode gets its own env (LiberoEnv is not thread-safe for reuse)
    envs = []
    for _ei in range(G):
        envs.append(LiberoEnv(libero_python=libero_python))
        if (_ei + 1) % 10 == 0 or _ei == G - 1:
            print(f"  envs created: {_ei+1}/{G}", flush=True)

    # Assign initial states: same-state grouping.
    # G episodes → G//group_size unique states, each repeated group_size times
    num_unique = max(1, G // group_size)
    _rng = np.random.RandomState(seed + group_idx)
    unique_states = _rng.randint(0, n_initial_states, size=num_unique)
    state_ids = np.repeat(unique_states, group_size)[:G]  # [s0,s0,s0,s0, s1,s1, ...]

    episodes = [None] * G
    try:
        with ThreadPoolExecutor(max_workers=n_workers) as pool:
            futures = {}
            for g in range(G):
                fut = pool.submit(
                    _action_token_rollout_one,
                    env=envs[g],
                    batch_server=batch_server,
                    suite_name=suite_name,
                    task_id=task_id,
                    state_idx=int(state_ids[g]),
                    action_norm_stats=action_norm_stats,
                    max_steps=max_steps,
                    chunk_len=chunk_len,
                    num_steps_wait=num_steps_wait,
                    seed=seed + g,
                    record_video=(video_dir is not None),
                    episode_idx=g,
                    group_idx=group_idx,
                    video_dir=video_dir,
                    store_images=store_images,
                    reward_coef=reward_coef,
                )
                futures[fut] = g
            done_count = 0
            success_count = 0
            for fut in as_completed(futures):
                g_idx = futures[fut]
                ep = fut.result()
                episodes[g_idx] = ep
                done_count += 1
                if ep.success:
                    success_count += 1
                print(f"  [rollout][dev={device}] ep {done_count}/{G} done "
                      f"({'SUCCESS' if ep.success else 'fail'}, "
                      f"{ep.env_steps} steps) "
                      f"[{success_count}/{done_count} success so far]", flush=True)
    finally:
        for env in envs:
            env.close()
        if _own_server:
            batch_server.stop()

    return episodes

def compute_action_token_gae(
    episode: ActionTokenEpisode,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
):
    """
    Compute GAE advantages and returns for a single RLActionToken episode.

    Returns:
        advantages: list of floats (len = finish_step)
        returns: list of floats (len = finish_step)
    """
    steps = episode.step_records
    n = episode.finish_step
    if n == 0:
        return [], []

    values = [s.value for s in steps[:n]]
    # Terminal value = 0 (episode ended)
    advantages = [0.0] * n
    returns = [0.0] * n

    # Sparse reward at last step only
    rewards = [0.0] * n
    rewards[-1] = episode.reward

    gae = 0.0
    for t in reversed(range(n)):
        next_value = values[t + 1] if t + 1 < n else 0.0
        delta = rewards[t] + gamma * next_value - values[t]
        gae = delta + gamma * gae_lambda * gae
        advantages[t] = gae
        returns[t] = gae + values[t]

    return advantages, returns