Upload B1K model checkpoints, training recipes, and evaluation artifacts (part 3)

884eb94 verified about 1 month ago

14 kB

"""B1K policy wrapper with action compression, rolling inpainting, and stage voting."""

import logging
import numpy as np
import torch
import dataclasses
from collections import deque

from openpi_client.base_policy import BasePolicy
from openpi_client.image_tools import resize_with_pad
from b1k.policies.b1k_policy import extract_state_from_proprio
from b1k.models.pi_behavior_config import TASK_NUM_STAGES
from b1k.shared.correction_rules import apply_correction_rules, check_gripper_variation
from b1k_eval_contract import PROPRIOCEPTION_INDICES

logger = logging.getLogger(__name__)

RESIZE_SIZE = 224


@dataclasses.dataclass
class B1KWrapperConfig:
    """Configuration for B1K policy wrapper execution parameters."""
    actions_to_execute: int = 26
    actions_to_keep: int = 4
    execute_in_n_steps: int = 20
    history_len: int = 3
    votes_to_promote: int = 2
    time_threshold_inpaint: float = 0.3
    num_steps: int = 20
    apply_eval_tricks: bool = True


class B1KPolicyWrapper():
    """B1K policy wrapper for PI_BEHAVIOR models with action compression, rolling inpainting, and stage voting."""
    
    def __init__(
        self, 
        policy: BasePolicy,
        text_prompt: str = "PI_BEHAVIOR model (task-conditioned)",  # Not used, kept for compatibility
        action_horizon: int = 30,
        task_id: int | None = None,
        config: B1KWrapperConfig = None,
        checkpoint_switcher = None,
    ) -> None:
        self.base_policy = policy
        self.policy = policy
        self.checkpoint_switcher = checkpoint_switcher
        self.text_prompt = text_prompt
        self.action_horizon = action_horizon
        self.config = config if config is not None else B1KWrapperConfig()
        
        # Validate configuration
        if self.config.actions_to_execute + self.config.actions_to_keep > self.action_horizon:
            raise ValueError(
                f"actions_to_execute + actions_to_keep exceeds action_horizon"
            )
        
        # PI_BEHAVIOR specific (always True for B1K)
        self.task_id = task_id
        self.current_stage = 0
        self.prediction_history = deque([], maxlen=self.config.history_len)
        
        # Control loop variables
        self.last_actions = None
        self.action_index = 0
        self.step_count = 0
        self.prediction_count = 0
        self.next_initial_actions = None
    
    def reset(self):
        """Reset policy state."""
        self.policy.reset()
        self.last_actions = None
        self.action_index = 0
        self.step_count = 0
        self.prediction_count = 0
        self.next_initial_actions = None
        self.current_stage = 0
        self.prediction_history.clear()
        logger.info(f"Policy reset - Task ID: {self.task_id}, Action horizon: {self.action_horizon}")
    
    def _handle_task_change(self, new_task_id):
        """Handle task ID change by switching checkpoint and resetting state."""
        if self.task_id != new_task_id:
            old_task_id = self.task_id
            self.task_id = new_task_id
            
            logger.info(f"🔄 Task change detected: {old_task_id} → {new_task_id} (max stages: {TASK_NUM_STAGES[new_task_id]})")
            
            if self.checkpoint_switcher:
                new_policy = self.checkpoint_switcher.get_policy_for_task(new_task_id)
                if new_policy is not self.policy:
                    logger.info(f"📦 Switching checkpoint: task {old_task_id} → {new_task_id}")
                    self.base_policy = new_policy
                    self.policy = new_policy
                    self.policy.reset()
            
            self.current_stage = 0
            self.prediction_history.clear()
            self.last_actions = None
            self.action_index = 0
            self.next_initial_actions = None

    def process_obs(self, obs: dict) -> dict:
        """Process observation to match model input format."""
        prop_state = obs["robot_r1::proprio"]
        
        head_original = obs["robot_r1::robot_r1:zed_link:Camera:0::rgb"][..., :3]
        left_original = obs["robot_r1::robot_r1:left_realsense_link:Camera:0::rgb"][..., :3]
        right_original = obs["robot_r1::robot_r1:right_realsense_link:Camera:0::rgb"][..., :3]
        
        # Resize images
        head_resized = resize_with_pad(head_original, RESIZE_SIZE, RESIZE_SIZE)
        left_resized = resize_with_pad(left_original, RESIZE_SIZE, RESIZE_SIZE)
        right_resized = resize_with_pad(right_original, RESIZE_SIZE, RESIZE_SIZE)
        
        return {
            "observation/egocentric_camera": head_resized,
            "observation/wrist_image_left": left_resized,
            "observation/wrist_image_right": right_resized,
            "observation/state": prop_state,
            "prompt": self.text_prompt,
        }
    
    def update_current_stage(self, predicted_subtask_logits):
        """Update current stage using majority voting."""
        if self.task_id is None:
            return
            
        max_stage = TASK_NUM_STAGES[self.task_id] - 1
        predicted_stage = int(np.argmax(predicted_subtask_logits))
        
        if predicted_stage > max_stage:
            predicted_stage = max_stage
        
        self.prediction_history.append(predicted_stage)
        
        if len(self.prediction_history) == self.config.history_len:
            next_stage = self.current_stage + 1
            
            if next_stage <= max_stage:
                votes_for_next = sum(1 for pred in self.prediction_history if pred == next_stage)
                votes_to_skip = sum(1 for pred in self.prediction_history if pred == next_stage + 1)
                votes_to_go_back = sum(1 for pred in self.prediction_history if pred == self.current_stage - 1)
                
                if votes_for_next >= self.config.votes_to_promote:
                    old_stage = self.current_stage
                    self.current_stage = next_stage
                    self.prediction_history.clear()
                    logger.info(f"⬆️  Stage advanced: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
                elif votes_to_skip == self.config.history_len:
                    old_stage = self.current_stage
                    self.current_stage = next_stage
                    self.prediction_history.clear()
                    logger.info(f"⏭️  Stage skipped: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
                elif votes_to_go_back == self.config.history_len and self.current_stage > 0:
                    old_stage = self.current_stage
                    self.current_stage -= 1
                    self.prediction_history.clear()
                    logger.info(f"⬅️  Stage went back: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
    
    def prepare_batch_for_pi_behavior(self, batch):
        """Prepare batch for PI_BEHAVIOR model by adding task_id and current_stage."""
        task_id = self.task_id if self.task_id is not None else -1
        batch_copy = batch.copy()
        if "prompt" in batch_copy:
            del batch_copy["prompt"]
        
        batch_copy["tokenized_prompt"] = np.array([task_id, self.current_stage], dtype=np.int32)
        batch_copy["tokenized_prompt_mask"] = np.array([True, True], dtype=bool)
        batch_copy["subtask_state"] = np.array(self.current_stage, dtype=np.int32)
        
        return batch_copy
    
    def _interpolate_actions(self, actions, target_steps):
        """Interpolate actions using cubic spline."""
        from scipy.interpolate import interp1d
        
        original_indices = np.linspace(0, len(actions)-1, len(actions))
        target_indices = np.linspace(0, len(actions)-1, target_steps)
        
        interpolated = np.zeros((target_steps, actions.shape[1]))
        for dim in range(actions.shape[1]):
            f = interp1d(original_indices, actions[:, dim], kind='cubic')
            interpolated[:, dim] = f(target_indices)
        
        return interpolated

    def act(self, obs: dict) -> torch.Tensor:
        """Main action function."""
        
        # Extract task_id from observations
        if "task_id" in obs:
            new_task_id = int(obs["task_id"][0])
            self._handle_task_change(new_task_id)
        
        raw_state = obs["robot_r1::proprio"]
        current_state = extract_state_from_proprio(raw_state)
        
        # Check if we need new actions
        if self.last_actions is None or self.action_index >= self.config.execute_in_n_steps:
            
            # Process observation
            model_input = self.process_obs(obs)
            model_input = self.prepare_batch_for_pi_behavior(model_input)
            
            # Add rolling inpainting if available
            if self.next_initial_actions is not None and ("initial_actions" not in model_input or model_input["initial_actions"] is None):
                model_input["initial_actions"] = self.next_initial_actions
            
            # Get prediction
            if "initial_actions" in model_input and model_input["initial_actions"] is not None:
                output = self.policy.infer(model_input, initial_actions=model_input["initial_actions"])
            else:
                output = self.policy.infer(model_input)
            
            actions = output["actions"]
            
            # Ensure correct shape
            if len(actions.shape) == 3:
                actions = actions[0]
            if actions.shape[1] > 23:
                actions = actions[:, :23]
            
            # Apply eval tricks if enabled
            should_compress = self.config.execute_in_n_steps < self.config.actions_to_execute
            
            if self.config.apply_eval_tricks:
                if self.task_id is not None:
                    actions_before = actions.copy()
                    actions, corrected_stage = apply_correction_rules(
                        self.task_id, self.current_stage, current_state, actions
                    )
                    
                    # Log if stage was corrected
                    if corrected_stage != self.current_stage:
                        logger.info(f"🔧 Correction rule: Stage corrected {self.current_stage} → {corrected_stage} (task {self.task_id}, step {self.step_count})")
                        self.current_stage = corrected_stage
                        self.prediction_history.clear()
                    
                    # Log if actions were modified
                    if not np.allclose(actions_before, actions, rtol=1e-3):
                        max_diff = np.max(np.abs(actions_before - actions))
                        logger.info(f"🔧 Correction rule: Actions modified (max diff: {max_diff:.4f}, task {self.task_id}, stage {self.current_stage})")
                
                if should_compress:
                    has_high_variation, mean_var, max_var = check_gripper_variation(
                        actions, self.config.actions_to_execute
                    )
                    if has_high_variation:
                        should_compress = False
                        logger.info(f"🔧 Gripper variation: Compression disabled (mean: {mean_var:.4f}, max: {max_var:.4f})")
            
            # Determine execution parameters
            actions_to_execute = self.config.actions_to_execute if should_compress else self.config.execute_in_n_steps
            execute_steps = self.config.execute_in_n_steps
            
            # Save actions for next inpainting (before compression)
            inpainting_start = actions_to_execute
            inpainting_end = inpainting_start + self.config.actions_to_keep
            
            if len(actions) >= inpainting_end:
                self.next_initial_actions = actions[inpainting_start:inpainting_end].copy()
            else:
                self.next_initial_actions = None
            
            # Extract and compress actions
            self.last_actions = actions[:actions_to_execute].copy()
            
            if should_compress:
                compressed_actions = self._interpolate_actions(self.last_actions, execute_steps)
                compression_factor = actions_to_execute / execute_steps
                compressed_actions[:, :3] *= compression_factor  # Scale velocities
                self.last_actions = compressed_actions
            
            self.action_index = 0
            self.prediction_count += 1
            
            # Log prediction details (at lower frequency, every 10 predictions)
            if self.prediction_count % 10 == 0:
                compression_status = f"compressed {actions_to_execute}→{execute_steps}" if should_compress else f"uncompressed ({execute_steps})"
                logger.info(f"🎯 Prediction #{self.prediction_count} | Actions: {compression_status} | Inpainting: {self.next_initial_actions is not None}")
            
            # Update stage based on model predictions
            if "subtask_logits" in output:
                self.update_current_stage(output["subtask_logits"])
        
        # Get current action from sequence
        if self.action_index >= len(self.last_actions):
            self.action_index = 0
            
        current_action = self.last_actions[self.action_index]
        self.action_index += 1
        self.step_count += 1
        
        # Log progress every 100 steps
        if self.step_count % 100 == 0:
            logger.info(f"📊 Step {self.step_count} | Task: {self.task_id} | Stage: {self.current_stage}/{TASK_NUM_STAGES[self.task_id]-1} | Predictions: {self.prediction_count}")
        
        # Convert to torch tensor
        action_tensor = torch.from_numpy(current_action).float()
        if len(action_tensor) > 23:
            action_tensor = action_tensor[:23]
        
        return action_tensor

	"""B1K policy wrapper with action compression, rolling inpainting, and stage voting."""

	import logging
	import numpy as np
	import torch
	import dataclasses
	from collections import deque

	from openpi_client.base_policy import BasePolicy
	from openpi_client.image_tools import resize_with_pad
	from b1k.policies.b1k_policy import extract_state_from_proprio
	from b1k.models.pi_behavior_config import TASK_NUM_STAGES
	from b1k.shared.correction_rules import apply_correction_rules, check_gripper_variation
	from b1k_eval_contract import PROPRIOCEPTION_INDICES

	logger = logging.getLogger(__name__)

	RESIZE_SIZE = 224


	@dataclasses.dataclass
	class B1KWrapperConfig:
	"""Configuration for B1K policy wrapper execution parameters."""
	actions_to_execute: int = 26
	actions_to_keep: int = 4
	execute_in_n_steps: int = 20
	history_len: int = 3
	votes_to_promote: int = 2
	time_threshold_inpaint: float = 0.3
	num_steps: int = 20
	apply_eval_tricks: bool = True


	class B1KPolicyWrapper():
	"""B1K policy wrapper for PI_BEHAVIOR models with action compression, rolling inpainting, and stage voting."""

	def __init__(
	self,
	policy: BasePolicy,
	text_prompt: str = "PI_BEHAVIOR model (task-conditioned)", # Not used, kept for compatibility
	action_horizon: int = 30,
	task_id: int \| None = None,
	config: B1KWrapperConfig = None,
	checkpoint_switcher = None,
	) -> None:
	self.base_policy = policy
	self.policy = policy
	self.checkpoint_switcher = checkpoint_switcher
	self.text_prompt = text_prompt
	self.action_horizon = action_horizon
	self.config = config if config is not None else B1KWrapperConfig()

	# Validate configuration
	if self.config.actions_to_execute + self.config.actions_to_keep > self.action_horizon:
	raise ValueError(
	f"actions_to_execute + actions_to_keep exceeds action_horizon"
	)

	# PI_BEHAVIOR specific (always True for B1K)
	self.task_id = task_id
	self.current_stage = 0
	self.prediction_history = deque([], maxlen=self.config.history_len)

	# Control loop variables
	self.last_actions = None
	self.action_index = 0
	self.step_count = 0
	self.prediction_count = 0
	self.next_initial_actions = None

	def reset(self):
	"""Reset policy state."""
	self.policy.reset()
	self.last_actions = None
	self.action_index = 0
	self.step_count = 0
	self.prediction_count = 0
	self.next_initial_actions = None
	self.current_stage = 0
	self.prediction_history.clear()
	logger.info(f"Policy reset - Task ID: {self.task_id}, Action horizon: {self.action_horizon}")

	def _handle_task_change(self, new_task_id):
	"""Handle task ID change by switching checkpoint and resetting state."""
	if self.task_id != new_task_id:
	old_task_id = self.task_id
	self.task_id = new_task_id

	logger.info(f"🔄 Task change detected: {old_task_id} → {new_task_id} (max stages: {TASK_NUM_STAGES[new_task_id]})")

	if self.checkpoint_switcher:
	new_policy = self.checkpoint_switcher.get_policy_for_task(new_task_id)
	if new_policy is not self.policy:
	logger.info(f"📦 Switching checkpoint: task {old_task_id} → {new_task_id}")
	self.base_policy = new_policy
	self.policy = new_policy
	self.policy.reset()

	self.current_stage = 0
	self.prediction_history.clear()
	self.last_actions = None
	self.action_index = 0
	self.next_initial_actions = None

	def process_obs(self, obs: dict) -> dict:
	"""Process observation to match model input format."""
	prop_state = obs["robot_r1::proprio"]

	head_original = obs["robot_r1::robot_r1:zed_link:Camera:0::rgb"][..., :3]
	left_original = obs["robot_r1::robot_r1:left_realsense_link:Camera:0::rgb"][..., :3]
	right_original = obs["robot_r1::robot_r1:right_realsense_link:Camera:0::rgb"][..., :3]

	# Resize images
	head_resized = resize_with_pad(head_original, RESIZE_SIZE, RESIZE_SIZE)
	left_resized = resize_with_pad(left_original, RESIZE_SIZE, RESIZE_SIZE)
	right_resized = resize_with_pad(right_original, RESIZE_SIZE, RESIZE_SIZE)

	return {
	"observation/egocentric_camera": head_resized,
	"observation/wrist_image_left": left_resized,
	"observation/wrist_image_right": right_resized,
	"observation/state": prop_state,
	"prompt": self.text_prompt,
	}

	def update_current_stage(self, predicted_subtask_logits):
	"""Update current stage using majority voting."""
	if self.task_id is None:
	return

	max_stage = TASK_NUM_STAGES[self.task_id] - 1
	predicted_stage = int(np.argmax(predicted_subtask_logits))

	if predicted_stage > max_stage:
	predicted_stage = max_stage

	self.prediction_history.append(predicted_stage)

	if len(self.prediction_history) == self.config.history_len:
	next_stage = self.current_stage + 1

	if next_stage <= max_stage:
	votes_for_next = sum(1 for pred in self.prediction_history if pred == next_stage)
	votes_to_skip = sum(1 for pred in self.prediction_history if pred == next_stage + 1)
	votes_to_go_back = sum(1 for pred in self.prediction_history if pred == self.current_stage - 1)

	if votes_for_next >= self.config.votes_to_promote:
	old_stage = self.current_stage
	self.current_stage = next_stage
	self.prediction_history.clear()
	logger.info(f"⬆️ Stage advanced: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
	elif votes_to_skip == self.config.history_len:
	old_stage = self.current_stage
	self.current_stage = next_stage
	self.prediction_history.clear()
	logger.info(f"⏭️ Stage skipped: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")
	elif votes_to_go_back == self.config.history_len and self.current_stage > 0:
	old_stage = self.current_stage
	self.current_stage -= 1
	self.prediction_history.clear()
	logger.info(f"⬅️ Stage went back: {old_stage} → {self.current_stage} (task {self.task_id}, step {self.step_count})")

	def prepare_batch_for_pi_behavior(self, batch):
	"""Prepare batch for PI_BEHAVIOR model by adding task_id and current_stage."""
	task_id = self.task_id if self.task_id is not None else -1
	batch_copy = batch.copy()
	if "prompt" in batch_copy:
	del batch_copy["prompt"]

	batch_copy["tokenized_prompt"] = np.array([task_id, self.current_stage], dtype=np.int32)
	batch_copy["tokenized_prompt_mask"] = np.array([True, True], dtype=bool)
	batch_copy["subtask_state"] = np.array(self.current_stage, dtype=np.int32)

	return batch_copy

	def _interpolate_actions(self, actions, target_steps):
	"""Interpolate actions using cubic spline."""
	from scipy.interpolate import interp1d

	original_indices = np.linspace(0, len(actions)-1, len(actions))
	target_indices = np.linspace(0, len(actions)-1, target_steps)

	interpolated = np.zeros((target_steps, actions.shape[1]))
	for dim in range(actions.shape[1]):
	f = interp1d(original_indices, actions[:, dim], kind='cubic')
	interpolated[:, dim] = f(target_indices)

	return interpolated

	def act(self, obs: dict) -> torch.Tensor:
	"""Main action function."""

	# Extract task_id from observations
	if "task_id" in obs:
	new_task_id = int(obs["task_id"][0])
	self._handle_task_change(new_task_id)

	raw_state = obs["robot_r1::proprio"]
	current_state = extract_state_from_proprio(raw_state)

	# Check if we need new actions
	if self.last_actions is None or self.action_index >= self.config.execute_in_n_steps:

	# Process observation
	model_input = self.process_obs(obs)
	model_input = self.prepare_batch_for_pi_behavior(model_input)

	# Add rolling inpainting if available
	if self.next_initial_actions is not None and ("initial_actions" not in model_input or model_input["initial_actions"] is None):
	model_input["initial_actions"] = self.next_initial_actions

	# Get prediction
	if "initial_actions" in model_input and model_input["initial_actions"] is not None:
	output = self.policy.infer(model_input, initial_actions=model_input["initial_actions"])
	else:
	output = self.policy.infer(model_input)

	actions = output["actions"]

	# Ensure correct shape
	if len(actions.shape) == 3:
	actions = actions[0]
	if actions.shape[1] > 23:
	actions = actions[:, :23]

	# Apply eval tricks if enabled
	should_compress = self.config.execute_in_n_steps < self.config.actions_to_execute

	if self.config.apply_eval_tricks:
	if self.task_id is not None:
	actions_before = actions.copy()
	actions, corrected_stage = apply_correction_rules(
	self.task_id, self.current_stage, current_state, actions
	)

	# Log if stage was corrected
	if corrected_stage != self.current_stage:
	logger.info(f"🔧 Correction rule: Stage corrected {self.current_stage} → {corrected_stage} (task {self.task_id}, step {self.step_count})")
	self.current_stage = corrected_stage
	self.prediction_history.clear()

	# Log if actions were modified
	if not np.allclose(actions_before, actions, rtol=1e-3):
	max_diff = np.max(np.abs(actions_before - actions))
	logger.info(f"🔧 Correction rule: Actions modified (max diff: {max_diff:.4f}, task {self.task_id}, stage {self.current_stage})")

	if should_compress:
	has_high_variation, mean_var, max_var = check_gripper_variation(
	actions, self.config.actions_to_execute
	)
	if has_high_variation:
	should_compress = False
	logger.info(f"🔧 Gripper variation: Compression disabled (mean: {mean_var:.4f}, max: {max_var:.4f})")

	# Determine execution parameters
	actions_to_execute = self.config.actions_to_execute if should_compress else self.config.execute_in_n_steps
	execute_steps = self.config.execute_in_n_steps

	# Save actions for next inpainting (before compression)
	inpainting_start = actions_to_execute
	inpainting_end = inpainting_start + self.config.actions_to_keep

	if len(actions) >= inpainting_end:
	self.next_initial_actions = actions[inpainting_start:inpainting_end].copy()
	else:
	self.next_initial_actions = None

	# Extract and compress actions
	self.last_actions = actions[:actions_to_execute].copy()

	if should_compress:
	compressed_actions = self._interpolate_actions(self.last_actions, execute_steps)
	compression_factor = actions_to_execute / execute_steps
	compressed_actions[:, :3] *= compression_factor # Scale velocities
	self.last_actions = compressed_actions

	self.action_index = 0
	self.prediction_count += 1

	# Log prediction details (at lower frequency, every 10 predictions)
	if self.prediction_count % 10 == 0:
	compression_status = f"compressed {actions_to_execute}→{execute_steps}" if should_compress else f"uncompressed ({execute_steps})"
	logger.info(f"🎯 Prediction #{self.prediction_count} \| Actions: {compression_status} \| Inpainting: {self.next_initial_actions is not None}")

	# Update stage based on model predictions
	if "subtask_logits" in output:
	self.update_current_stage(output["subtask_logits"])

	# Get current action from sequence
	if self.action_index >= len(self.last_actions):
	self.action_index = 0

	current_action = self.last_actions[self.action_index]
	self.action_index += 1
	self.step_count += 1

	# Log progress every 100 steps
	if self.step_count % 100 == 0:
	logger.info(f"📊 Step {self.step_count} \| Task: {self.task_id} \| Stage: {self.current_stage}/{TASK_NUM_STAGES[self.task_id]-1} \| Predictions: {self.prediction_count}")

	# Convert to torch tensor
	action_tensor = torch.from_numpy(current_action).float()
	if len(action_tensor) > 23:
	action_tensor = action_tensor[:23]

	return action_tensor