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

884eb94 verified about 1 month ago

13.3 kB

"""Data transforms for BEHAVIOR-1K dataset.
Standard transforms imported from OpenPI.
B1K-specific: TaskIndexToTaskId, ComputeSubtaskStateFromMeta, TokenizeFASTActions
Reference: https://github.com/wensi-ai/openpi/tree/behavior
"""

import dataclasses
import logging
import numpy as np

# Import all standard transforms from OpenPI
from openpi.transforms import (
    DataTransformFn,
    DataDict,
    Group,
    CompositeTransform,
    compose,
    RepackTransform,
    Normalize,
    Unnormalize,
    ResizeImages,
    SubsampleActions,
    DeltaActions,
    AbsoluteActions,
    PadStatesAndActions,
    PromptFromLeRobotTask,  # Not used for PI_BEHAVIOR but kept for compatibility
    flatten_dict,
    unflatten_dict,
    transform_dict,
    apply_tree,
    pad_to_dim,
    make_bool_mask,
)

from b1k.models.pi_behavior_config import TASK_NUM_STAGES
from b1k.shared.normalize import NormStats


@dataclasses.dataclass(frozen=True)
class TaskIndexToTaskId(DataTransformFn):
    """Directly converts task_index to task_id for PI_BEHAVIOR model.
    
    PI_BEHAVIOR uses task embeddings instead of text prompts. This transform
    converts the dataset's task_index to task_id and prepares tokenized_prompt
    as [task_id, subtask_state] for the model.
    
    Assumes:
    - data["task_index"] exists (from dataset)
    - data["subtask_state"] exists (computed by ComputeSubtaskStateFromMeta)
    
    Creates:
    - data["tokenized_prompt"]: np.array([task_id, subtask_state], dtype=int32)
    - data["tokenized_prompt_mask"]: np.array([True, True], dtype=bool)
    """
    
    # Optional task remapping (dataset task_index → model task_id)
    # If None, assumes direct mapping (task_index == task_id)
    task_mapping: dict[int, int] | None = None
    
    def __call__(self, data: DataDict) -> DataDict:
        # During inference, task_id might be provided directly instead of task_index
        if "task_id" in data:
            # Direct task_id provided (inference mode)
            task_id = int(data["task_id"])
        elif "task_index" in data:
            # task_index provided (training mode)
            task_index = int(data["task_index"])
            
            # Apply mapping if provided, otherwise use task_index directly as task_id
            if self.task_mapping is not None:
                if task_index not in self.task_mapping:
                    raise ValueError(f"task_index {task_index} not found in mapping")
                task_id = self.task_mapping[task_index]
            else:
                task_id = task_index
        else:
            # During inference, if neither is provided, this transform should be skipped
            # The tokenized_prompt should already be set up correctly
            if "tokenized_prompt" in data:
                return data  # Already has tokenized_prompt, skip this transform
            raise ValueError("Either task_index, task_id, or tokenized_prompt is required for PI_BEHAVIOR model")

        # Pack task_id and subtask_state (if available) into tokenized_prompt
        if "subtask_state" in data:
            subtask_state = int(data["subtask_state"])
            prompt_tokens = np.array([task_id, subtask_state], dtype=np.int32)  # [task_id, subtask_state]
            prompt_mask = np.array([True, True], dtype=bool)
        else:
            prompt_tokens = np.array([task_id], dtype=np.int32)  # Just [task_id]
            prompt_mask = np.array([True], dtype=bool)

        return {
            **data, 
            "tokenized_prompt": prompt_tokens,
            "tokenized_prompt_mask": prompt_mask
        }


@dataclasses.dataclass(frozen=True)
class ComputeSubtaskStateFromMeta(DataTransformFn):
    """Computes subtask state from timestamp using dataset.meta.episodes.
    
    Divides episode into task-specific number of stages based on episode length.
    Requires dataset reference to access episode metadata.
    
    Args:
        dataset: Dataset instance with meta.episodes containing episode_length
    
    Assumes:
    - data["episode_index"] exists
    - data["timestamp"] exists (frame index within episode)
    - data["task_index"] exists (for task-specific stage count)
    - dataset.meta.episodes contains episode_length for each episode
    
    Creates:
    - data["subtask_state"]: Stage index (0 to num_stages-1)
    """
    
    dataset: object | None = None  # Will be set by data loader
    
    def __call__(self, data: DataDict) -> DataDict:
        if self.dataset is None:
            # During inference or when dataset is not available, use default
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
        
        if "episode_index" not in data or "timestamp" not in data or "task_index" not in data:
            # Missing required fields, default to stage 0
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
        
        episode_index = int(data["episode_index"])
        timestamp = float(data["timestamp"])
        task_index = int(data["task_index"])
        
        # Validate task_index
        if not (0 <= task_index < 50):
            logging.warning(f"Invalid task_index {task_index}, using stage 0")
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
        
        # Get number of stages for this task
        num_stages = TASK_NUM_STAGES[task_index]
        
        # Get episode length from dataset metadata
        if not hasattr(self.dataset, 'meta') or not hasattr(self.dataset.meta, 'episodes'):
            # No metadata available
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
            
        meta_episodes = self.dataset.meta.episodes
        
        if episode_index not in meta_episodes:
            logging.warning(f"Episode {episode_index} not found in metadata, using stage 0")
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
            
        episode_info = meta_episodes[episode_index]
        
        # Try 'length' first (standard key), then 'episode_length' (alternative)
        episode_length = episode_info.get('length', episode_info.get('episode_length', None))
        
        if episode_length is None or episode_length <= 0:
            logging.warning(f"Invalid episode_length for episode {episode_index}, using stage 0")
            data["subtask_state"] = np.array(0, dtype=np.int32)
            return data
        
        episode_length = float(episode_length)
        
        # CRITICAL: Convert timestamp (in seconds) to frames (30 FPS)
        # Dataset provides timestamp in seconds, episode_length is in frames
        current_step = timestamp * 30.0
        
        # Divide episode into num_stages equal parts
        frames_per_stage = episode_length / num_stages
        
        # Compute current stage (0-indexed)
        subtask_state = int(current_step / frames_per_stage)
        
        # Clamp to valid range [0, num_stages-1]
        subtask_state = max(0, min(subtask_state, num_stages - 1))
        
        data["subtask_state"] = np.array(subtask_state, dtype=np.int32)
        return data


@dataclasses.dataclass(frozen=True)
class TokenizeFASTActions(DataTransformFn):
    """Tokenize actions to FAST discrete tokens for PI_BEHAVIOR auxiliary training.
    
    Converts continuous actions to discrete tokens using a trained FAST tokenizer.
    This is used for auxiliary training to improve action quality.
    
    Note: Uses GLOBAL normalization even if per-timestamp normalization is enabled
    elsewhere, to preserve temporal smoothness for better DCT compression.
    
    Args:
        tokenizer_path: Path to trained FAST tokenizer directory
        encoded_dim_ranges: List of (start, end) tuples for dimensions to encode
        max_fast_tokens: Maximum number of tokens (truncate/pad)
        norm_stats: Normalization stats for denorm/renorm if per-timestamp is used
        use_per_timestamp: Whether per-timestamp normalization is being used
    
    Assumes:
    - data["actions"] exists and is normalized
    - data["state"] exists (for delta computation)
    
    Creates:
    - data["fast_tokens"]: Discrete token indices [max_fast_tokens]
    - data["fast_token_mask"]: Boolean mask for valid tokens [max_fast_tokens]
    """
    
    # Path to FAST tokenizer
    tokenizer_path: str
    # Action dimensions to encode: [(start, end), ...]
    encoded_dim_ranges: list[tuple[int, int]]
    # Max tokens to predict
    max_fast_tokens: int = 180
    # Normalization stats (for denorm/renorm if per-timestamp is used)
    norm_stats: dict[str, NormStats] | None = None
    # Whether per-timestamp normalization is being used
    use_per_timestamp: bool = False
    
    def _get_tokenizer(self):
        """Lazy load tokenizer (called in each worker process)."""
        # Check if already loaded in this process
        if not hasattr(self, '_tokenizer'):
            import importlib.util
            from pathlib import Path
            
            tokenizer_dir = Path(self.tokenizer_path)
            processor_file = tokenizer_dir / "processing_action_tokenizer.py"
            
            if not processor_file.exists():
                raise FileNotFoundError(f"Tokenizer processing file not found at {processor_file}")
            
            # Load the module dynamically
            spec = importlib.util.spec_from_file_location("processing_action_tokenizer", processor_file)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            
            # Get the UniversalActionProcessor class
            UniversalActionProcessor = module.UniversalActionProcessor
            
            # Load the tokenizer using from_pretrained
            tokenizer = UniversalActionProcessor.from_pretrained(str(tokenizer_dir))
            
            # Store in this object (won't be pickled due to lazy loading pattern)
            object.__setattr__(self, '_tokenizer', tokenizer)
        
        return self._tokenizer
    
    def __call__(self, data: DataDict) -> DataDict:
        # Only tokenize if actions are present (training time)
        if "actions" not in data:
            return data
        
        actions = data["actions"]  # [H, D] - may be per-timestamp normalized
        
        # Extract encoded dimensions
        encoded_chunks = []
        for start, end in self.encoded_dim_ranges:
            encoded_chunks.append(actions[:, start:end])
        encoded_actions = np.concatenate(encoded_chunks, axis=-1)  # [H, D_encoded]
        
        # Convert to global QUANTILE normalization for FAST (always, regardless of per_timestamp)
        # This clips outliers and provides consistent [-1, 1] range for DCT compression
        if self.norm_stats is not None:
            action_stats = self.norm_stats['actions']
            
            # Extract stats for encoded dimensions only
            # Build dimension indices from ranges
            encoded_dims = []
            for start, end in self.encoded_dim_ranges:
                encoded_dims.extend(range(start, end))
            encoded_dims = np.array(encoded_dims)
            
            # Denormalize to get back to raw delta actions
            if self.use_per_timestamp:
                # Denormalize from per-timestamp MEAN/STD normalization
                per_ts_mean = action_stats.per_timestamp_mean[:, encoded_dims]  # [H, D_encoded]
                per_ts_std = action_stats.per_timestamp_std[:, encoded_dims]    # [H, D_encoded]
                actions_delta = encoded_actions * per_ts_std + per_ts_mean
            else:
                # Denormalize from global MEAN/STD normalization
                global_mean = action_stats.mean[encoded_dims]  # [D_encoded]
                global_std = action_stats.std[encoded_dims]    # [D_encoded]
                actions_delta = encoded_actions * global_std + global_mean
            
            # Renormalize using GLOBAL QUANTILE normalization (for FAST)
            # This clips to [q01, q99] and maps to [-1, 1]
            q01 = action_stats.q01[encoded_dims]  # [D_encoded]
            q99 = action_stats.q99[encoded_dims]  # [D_encoded]
            actions_delta = np.clip(actions_delta, q01, q99)
            encoded_actions = 2.0 * (actions_delta - q01) / np.maximum(q99 - q01, 1e-6) - 1.0
        
        # Tokenize
        tokenizer = self._get_tokenizer()
        tokens = tokenizer(encoded_actions[None])[0]
        
        if isinstance(tokens, list):
            tokens = np.array(tokens)
        
        # Truncate or pad to max_fast_tokens
        if len(tokens) > self.max_fast_tokens:
            tokens = tokens[:self.max_fast_tokens]
            mask = np.ones(self.max_fast_tokens, dtype=bool)
        else:
            mask = np.concatenate([
                np.ones(len(tokens), dtype=bool),
                np.zeros(self.max_fast_tokens - len(tokens), dtype=bool)
            ])
            tokens = np.pad(tokens, (0, self.max_fast_tokens - len(tokens)), constant_values=0)
        
        data["fast_tokens"] = tokens.astype(np.int32)
        data["fast_token_mask"] = mask
        
        return data

	"""Data transforms for BEHAVIOR-1K dataset.

	Standard transforms imported from OpenPI.
	B1K-specific: TaskIndexToTaskId, ComputeSubtaskStateFromMeta, TokenizeFASTActions

	Reference: https://github.com/wensi-ai/openpi/tree/behavior
	"""

	import dataclasses
	import logging
	import numpy as np

	# Import all standard transforms from OpenPI
	from openpi.transforms import (
	DataTransformFn,
	DataDict,
	Group,
	CompositeTransform,
	compose,
	RepackTransform,
	Normalize,
	Unnormalize,
	ResizeImages,
	SubsampleActions,
	DeltaActions,
	AbsoluteActions,
	PadStatesAndActions,
	PromptFromLeRobotTask, # Not used for PI_BEHAVIOR but kept for compatibility
	flatten_dict,
	unflatten_dict,
	transform_dict,
	apply_tree,
	pad_to_dim,
	make_bool_mask,
	)

	from b1k.models.pi_behavior_config import TASK_NUM_STAGES
	from b1k.shared.normalize import NormStats


	@dataclasses.dataclass(frozen=True)
	class TaskIndexToTaskId(DataTransformFn):
	"""Directly converts task_index to task_id for PI_BEHAVIOR model.

	PI_BEHAVIOR uses task embeddings instead of text prompts. This transform
	converts the dataset's task_index to task_id and prepares tokenized_prompt
	as [task_id, subtask_state] for the model.

	Assumes:
	- data["task_index"] exists (from dataset)
	- data["subtask_state"] exists (computed by ComputeSubtaskStateFromMeta)

	Creates:
	- data["tokenized_prompt"]: np.array([task_id, subtask_state], dtype=int32)
	- data["tokenized_prompt_mask"]: np.array([True, True], dtype=bool)
	"""

	# Optional task remapping (dataset task_index → model task_id)
	# If None, assumes direct mapping (task_index == task_id)
	task_mapping: dict[int, int] \| None = None

	def __call__(self, data: DataDict) -> DataDict:
	# During inference, task_id might be provided directly instead of task_index
	if "task_id" in data:
	# Direct task_id provided (inference mode)
	task_id = int(data["task_id"])
	elif "task_index" in data:
	# task_index provided (training mode)
	task_index = int(data["task_index"])

	# Apply mapping if provided, otherwise use task_index directly as task_id
	if self.task_mapping is not None:
	if task_index not in self.task_mapping:
	raise ValueError(f"task_index {task_index} not found in mapping")
	task_id = self.task_mapping[task_index]
	else:
	task_id = task_index
	else:
	# During inference, if neither is provided, this transform should be skipped
	# The tokenized_prompt should already be set up correctly
	if "tokenized_prompt" in data:
	return data # Already has tokenized_prompt, skip this transform
	raise ValueError("Either task_index, task_id, or tokenized_prompt is required for PI_BEHAVIOR model")

	# Pack task_id and subtask_state (if available) into tokenized_prompt
	if "subtask_state" in data:
	subtask_state = int(data["subtask_state"])
	prompt_tokens = np.array([task_id, subtask_state], dtype=np.int32) # [task_id, subtask_state]
	prompt_mask = np.array([True, True], dtype=bool)
	else:
	prompt_tokens = np.array([task_id], dtype=np.int32) # Just [task_id]
	prompt_mask = np.array([True], dtype=bool)

	return {
	**data,
	"tokenized_prompt": prompt_tokens,
	"tokenized_prompt_mask": prompt_mask
	}


	@dataclasses.dataclass(frozen=True)
	class ComputeSubtaskStateFromMeta(DataTransformFn):
	"""Computes subtask state from timestamp using dataset.meta.episodes.

	Divides episode into task-specific number of stages based on episode length.
	Requires dataset reference to access episode metadata.

	Args:
	dataset: Dataset instance with meta.episodes containing episode_length

	Assumes:
	- data["episode_index"] exists
	- data["timestamp"] exists (frame index within episode)
	- data["task_index"] exists (for task-specific stage count)
	- dataset.meta.episodes contains episode_length for each episode

	Creates:
	- data["subtask_state"]: Stage index (0 to num_stages-1)
	"""

	dataset: object \| None = None # Will be set by data loader

	def __call__(self, data: DataDict) -> DataDict:
	if self.dataset is None:
	# During inference or when dataset is not available, use default
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	if "episode_index" not in data or "timestamp" not in data or "task_index" not in data:
	# Missing required fields, default to stage 0
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	episode_index = int(data["episode_index"])
	timestamp = float(data["timestamp"])
	task_index = int(data["task_index"])

	# Validate task_index
	if not (0 <= task_index < 50):
	logging.warning(f"Invalid task_index {task_index}, using stage 0")
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	# Get number of stages for this task
	num_stages = TASK_NUM_STAGES[task_index]

	# Get episode length from dataset metadata
	if not hasattr(self.dataset, 'meta') or not hasattr(self.dataset.meta, 'episodes'):
	# No metadata available
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	meta_episodes = self.dataset.meta.episodes

	if episode_index not in meta_episodes:
	logging.warning(f"Episode {episode_index} not found in metadata, using stage 0")
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	episode_info = meta_episodes[episode_index]

	# Try 'length' first (standard key), then 'episode_length' (alternative)
	episode_length = episode_info.get('length', episode_info.get('episode_length', None))

	if episode_length is None or episode_length <= 0:
	logging.warning(f"Invalid episode_length for episode {episode_index}, using stage 0")
	data["subtask_state"] = np.array(0, dtype=np.int32)
	return data

	episode_length = float(episode_length)

	# CRITICAL: Convert timestamp (in seconds) to frames (30 FPS)
	# Dataset provides timestamp in seconds, episode_length is in frames
	current_step = timestamp * 30.0

	# Divide episode into num_stages equal parts
	frames_per_stage = episode_length / num_stages

	# Compute current stage (0-indexed)
	subtask_state = int(current_step / frames_per_stage)

	# Clamp to valid range [0, num_stages-1]
	subtask_state = max(0, min(subtask_state, num_stages - 1))

	data["subtask_state"] = np.array(subtask_state, dtype=np.int32)
	return data


	@dataclasses.dataclass(frozen=True)
	class TokenizeFASTActions(DataTransformFn):
	"""Tokenize actions to FAST discrete tokens for PI_BEHAVIOR auxiliary training.

	Converts continuous actions to discrete tokens using a trained FAST tokenizer.
	This is used for auxiliary training to improve action quality.

	Note: Uses GLOBAL normalization even if per-timestamp normalization is enabled
	elsewhere, to preserve temporal smoothness for better DCT compression.

	Args:
	tokenizer_path: Path to trained FAST tokenizer directory
	encoded_dim_ranges: List of (start, end) tuples for dimensions to encode
	max_fast_tokens: Maximum number of tokens (truncate/pad)
	norm_stats: Normalization stats for denorm/renorm if per-timestamp is used
	use_per_timestamp: Whether per-timestamp normalization is being used

	Assumes:
	- data["actions"] exists and is normalized
	- data["state"] exists (for delta computation)

	Creates:
	- data["fast_tokens"]: Discrete token indices [max_fast_tokens]
	- data["fast_token_mask"]: Boolean mask for valid tokens [max_fast_tokens]
	"""

	# Path to FAST tokenizer
	tokenizer_path: str
	# Action dimensions to encode: [(start, end), ...]
	encoded_dim_ranges: list[tuple[int, int]]
	# Max tokens to predict
	max_fast_tokens: int = 180
	# Normalization stats (for denorm/renorm if per-timestamp is used)
	norm_stats: dict[str, NormStats] \| None = None
	# Whether per-timestamp normalization is being used
	use_per_timestamp: bool = False

	def _get_tokenizer(self):
	"""Lazy load tokenizer (called in each worker process)."""
	# Check if already loaded in this process
	if not hasattr(self, '_tokenizer'):
	import importlib.util
	from pathlib import Path

	tokenizer_dir = Path(self.tokenizer_path)
	processor_file = tokenizer_dir / "processing_action_tokenizer.py"

	if not processor_file.exists():
	raise FileNotFoundError(f"Tokenizer processing file not found at {processor_file}")

	# Load the module dynamically
	spec = importlib.util.spec_from_file_location("processing_action_tokenizer", processor_file)
	module = importlib.util.module_from_spec(spec)
	spec.loader.exec_module(module)

	# Get the UniversalActionProcessor class
	UniversalActionProcessor = module.UniversalActionProcessor

	# Load the tokenizer using from_pretrained
	tokenizer = UniversalActionProcessor.from_pretrained(str(tokenizer_dir))

	# Store in this object (won't be pickled due to lazy loading pattern)
	object.__setattr__(self, '_tokenizer', tokenizer)

	return self._tokenizer

	def __call__(self, data: DataDict) -> DataDict:
	# Only tokenize if actions are present (training time)
	if "actions" not in data:
	return data

	actions = data["actions"] # [H, D] - may be per-timestamp normalized

	# Extract encoded dimensions
	encoded_chunks = []
	for start, end in self.encoded_dim_ranges:
	encoded_chunks.append(actions[:, start:end])
	encoded_actions = np.concatenate(encoded_chunks, axis=-1) # [H, D_encoded]

	# Convert to global QUANTILE normalization for FAST (always, regardless of per_timestamp)
	# This clips outliers and provides consistent [-1, 1] range for DCT compression
	if self.norm_stats is not None:
	action_stats = self.norm_stats['actions']

	# Extract stats for encoded dimensions only
	# Build dimension indices from ranges
	encoded_dims = []
	for start, end in self.encoded_dim_ranges:
	encoded_dims.extend(range(start, end))
	encoded_dims = np.array(encoded_dims)

	# Denormalize to get back to raw delta actions
	if self.use_per_timestamp:
	# Denormalize from per-timestamp MEAN/STD normalization
	per_ts_mean = action_stats.per_timestamp_mean[:, encoded_dims] # [H, D_encoded]
	per_ts_std = action_stats.per_timestamp_std[:, encoded_dims] # [H, D_encoded]
	actions_delta = encoded_actions * per_ts_std + per_ts_mean
	else:
	# Denormalize from global MEAN/STD normalization
	global_mean = action_stats.mean[encoded_dims] # [D_encoded]
	global_std = action_stats.std[encoded_dims] # [D_encoded]
	actions_delta = encoded_actions * global_std + global_mean

	# Renormalize using GLOBAL QUANTILE normalization (for FAST)
	# This clips to [q01, q99] and maps to [-1, 1]
	q01 = action_stats.q01[encoded_dims] # [D_encoded]
	q99 = action_stats.q99[encoded_dims] # [D_encoded]
	actions_delta = np.clip(actions_delta, q01, q99)
	encoded_actions = 2.0 * (actions_delta - q01) / np.maximum(q99 - q01, 1e-6) - 1.0

	# Tokenize
	tokenizer = self._get_tokenizer()
	tokens = tokenizer(encoded_actions[None])[0]

	if isinstance(tokens, list):
	tokens = np.array(tokens)

	# Truncate or pad to max_fast_tokens
	if len(tokens) > self.max_fast_tokens:
	tokens = tokens[:self.max_fast_tokens]
	mask = np.ones(self.max_fast_tokens, dtype=bool)
	else:
	mask = np.concatenate([
	np.ones(len(tokens), dtype=bool),
	np.zeros(self.max_fast_tokens - len(tokens), dtype=bool)
	])
	tokens = np.pad(tokens, (0, self.max_fast_tokens - len(tokens)), constant_values=0)

	data["fast_tokens"] = tokens.astype(np.int32)
	data["fast_token_mask"] = mask

	return data