Archive b1k-solution-code-metadata d12eb4b63fc6

e35ec21 verified 21 days ago

12.1 kB

"""
Evaluation tricks for improving success rates on BEHAVIOR-1K tasks.
This module contains:
1. Correction rules: Task-specific problem resolution (stage reset, gripper recovery)
2. Gripper variation check: Disable compression when grippers have large movement
Enable via --apply-eval-tricks flag (default: False)
"""

import logging
import numpy as np
from typing import Tuple, Optional

logger = logging.getLogger(__name__)


# ============================================================================
# GLOBAL CONSTANTS
# ============================================================================

# Gripper thresholds (normalized [-1, 1] range)
# Grippers are normalized from raw [0, 0.1] to [-1, 1]:
#   -1.0 = fully closed (raw 0.0)
#   -0.9 = mostly closed (raw 0.005)
#    0.0 = half open (raw 0.05)
#    0.9 = mostly open (raw 0.095)
#    1.0 = fully open (raw 0.1)
OPEN_THRESHOLD = 0.90
CLOSED_THRESHOLD = -0.98 # very strict threshold to avoid false positives

# Gripper variation threshold for compression check
GRIPPER_VARIATION_THRESHOLD = 0.2

# State indices
LEFT_GRIPPER_IDX = 14
RIGHT_GRIPPER_IDX = 22
LEFT_ARM_START_IDX = 7   # Left arm: state[7:14]
RIGHT_ARM_START_IDX = 15  # Right arm: state[15:22]
LEFT_ARM_END_IDX = 14
RIGHT_ARM_END_IDX = 22

# Base positions for arms (7 joints each)
BASE_POSITION_LEFT_ARM = np.array([-0.2, 0.0, 0.1, -1.4, 0.1, 0.7, 0.0])
BASE_POSITION_RIGHT_ARM = np.array([-0.2, 0.0, 0.1, -1.4, 0.1, 0.7, 0.0])


# ============================================================================
# GRIPPER CORRECTION POLICY
# ============================================================================

# Tasks where grippers should ALWAYS be open (any closure is an error)
ALWAYS_OPEN_LEFT_GRIPPER_TASKS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 42, 43, 44, 45, 47, 48}
ALWAYS_OPEN_RIGHT_GRIPPER_TASKS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 34, 35, 36, 37, 42, 43, 44, 47, 48, 49}

# Tasks where closure is not allowed BEFORE specified stage
# Format: {task_id: {'left': min_stage, 'right': min_stage}}
MIN_STAGE_FOR_CLOSURE = {
    0: {'left': 2, 'right': 2},      # Radio: allow closing from stage 2+
    30: {'right': 6},                 # setting_the_fire
    31: {'right': 8},                 # clean_boxing_gloves
    32: {'right': 5},                 # wash_a_baseball_cap
    33: {'right': 11},                # wash_dog_toys
    40: {'right': 4},                 # make_microwave_popcorn
    41: {'left': 14, 'right': 14},   # cook_cabbage
    45: {'right': 10},                # cook_hot_dogs
    46: {'left': 8, 'right': 8},     # cook_bacon
    49: {'left': 14},                 # make_pizza
}

# Tasks where right gripper is ALWAYS allowed to close (exceptions)
RIGHT_GRIPPER_ALWAYS_ALLOWED = {38, 39}  # spraying_for_bugs, spraying_fruit_trees


# ============================================================================
# CORRECTION RULES - ADD NEW RULES HERE
# ============================================================================

def task0_stage4_reset_to_stage2(
    task_id: int,
    stage: int,
    state: np.ndarray,
    actions: np.ndarray
) -> Optional[Tuple[np.ndarray, int]]:
    # Probably not very important for the final solution
    """
    Task 0 (Radio), Stage 4: Reset to stage 2 with arm/gripper recovery.
    
    PRIMARY: task_id=0 AND stage=4 → reset to stage 2 (always)
    
    SECONDARY (conditional):
    - If both grippers are fully open or fully closed:
      → Open closed grippers
    
    Args:
        task_id: Current task ID
        stage: Current stage
        state: Current state (23-dim)
        actions: Predicted actions (horizon, 23)
    
    Returns:
        (corrected_actions, corrected_stage) if rule applies, None otherwise
    """
    # Check primary condition
    if task_id != 0 or stage < 2:
        return None
    
    # PRIMARY: Always reset stage
    if stage == 4:
        logger.info(f"🔧 CORRECTION RULE: Task 0 (Radio), Stage 4 → Stage 2")
        corrected_stage = 2
    else:
        corrected_stage = stage

    # SECONDARY: Check gripper states
    left_gripper = state[LEFT_GRIPPER_IDX]
    right_gripper = state[RIGHT_GRIPPER_IDX]
    
    left_closed = left_gripper < CLOSED_THRESHOLD
    right_closed = right_gripper < CLOSED_THRESHOLD
    left_open = left_gripper > OPEN_THRESHOLD
    right_open = right_gripper > OPEN_THRESHOLD
    
    # Check if any gripper is in middle position (not fully open or closed)
    left_middle = not (left_open or left_closed)
    right_middle = not (right_open or right_closed)
    
    change_action = False

    # Copy current state to all actions as baseline
    corrected_actions = np.tile(state, (actions.shape[0], 1))

    if left_closed and not right_middle:
        corrected_actions[:, LEFT_GRIPPER_IDX] = 1.0
        change_action = True
        logger.info(f"   - Opening left gripper → 1.0")
    if right_closed and not left_middle:
        corrected_actions[:, RIGHT_GRIPPER_IDX] = 1.0
        change_action = True
        logger.info(f"   - Opening right gripper → 1.0")

    if not change_action:
        if stage == corrected_stage:
            return None

        corrected_actions = actions
    
    return corrected_actions, corrected_stage


def general_gripper_correction(
    task_id: int,
    stage: int,
    state: np.ndarray,
    actions: np.ndarray
) -> Optional[Tuple[np.ndarray, int]]:
    """
    General gripper correction rule for all tasks.
    
    Applies task-specific policies:
    1. Always-open tasks: Open gripper if closed (any stage)
    2. Stage-restricted tasks: Open gripper if closed before allowed stage
    3. Exceptions: Skip tasks where gripper is intentionally closed (e.g., spray bottles)
    
    Works in conjunction with task-specific rules:
    - Task 0: This rule handles stages ≤1, task-specific rule handles stage 4
    
    Args:
        task_id: Current task ID
        stage: Current stage
        state: np.ndarray,
        actions: Predicted actions (horizon, 23)
    
    Returns:
        (corrected_actions, corrected_stage) if rule applies, None otherwise
    """
    # Get gripper states
    left_gripper = state[LEFT_GRIPPER_IDX]
    right_gripper = state[RIGHT_GRIPPER_IDX]
    
    left_closed = left_gripper < CLOSED_THRESHOLD
    right_closed = right_gripper < CLOSED_THRESHOLD
    
    # Determine if correction is needed
    left_needs_opening = False
    right_needs_opening = False
    
    # Check left gripper
    if left_closed:
        # Always open if task requires it
        if task_id in ALWAYS_OPEN_LEFT_GRIPPER_TASKS:
            left_needs_opening = True
        # Check stage restriction
        elif task_id in MIN_STAGE_FOR_CLOSURE:
            min_stage_left = MIN_STAGE_FOR_CLOSURE[task_id].get('left')
            if min_stage_left is not None and stage < min_stage_left:
                left_needs_opening = True
    
    # Check right gripper
    if right_closed:
        # Exception: Always allow for spray bottle tasks
        if task_id in RIGHT_GRIPPER_ALWAYS_ALLOWED:
            pass  # Never correct these tasks
        # Always open if task requires it
        elif task_id in ALWAYS_OPEN_RIGHT_GRIPPER_TASKS:
            right_needs_opening = True
        # Check stage restriction
        elif task_id in MIN_STAGE_FOR_CLOSURE:
            min_stage_right = MIN_STAGE_FOR_CLOSURE[task_id].get('right')
            if min_stage_right is not None and stage < min_stage_right:
                right_needs_opening = True
    
    # Apply correction if needed
    if left_needs_opening or right_needs_opening:
        logger.info(f"🔧 GENERAL GRIPPER CORRECTION: Task {task_id}, Stage {stage}")
        
        # Copy current state to all actions
        corrected_actions = np.tile(state, (actions.shape[0], 1))
        
        if left_needs_opening:
            logger.info(f"   Left gripper closed ({left_gripper:.4f}) → Opening to 1.0")
            corrected_actions[:, LEFT_GRIPPER_IDX] = 1.0
        
        if right_needs_opening:
            logger.info(f"   Right gripper closed ({right_gripper:.4f}) → Opening to 1.0")
            corrected_actions[:, RIGHT_GRIPPER_IDX] = 1.0
        
        # Keep stage unchanged
        return corrected_actions, stage
    
    return None  # No correction needed


# ============================================================================
# RULE REGISTRY - ORGANIZE BY TASK_ID
# ============================================================================

# Task-specific rules (these run BEFORE the general rule)
# Dictionary mapping task_id -> list of rule functions
TASK_SPECIFIC_RULES = {
    0: [task0_stage4_reset_to_stage2],  # Radio has special stage 4 handling
    # Add more task-specific rules here as needed
}

# General rules that apply to ALL tasks (run after task-specific rules)
GENERAL_RULES = [
    general_gripper_correction,  # Applies to all tasks with gripper policies
]


# ============================================================================
# RULE APPLICATION
# ============================================================================

def apply_correction_rules(
    task_id: int,
    stage: int,
    state: np.ndarray,
    actions: np.ndarray
) -> Tuple[np.ndarray, int]:
    """
    Apply correction rules for the given task_id.
    
    Rules are applied in order until one matches:
    1. Task-specific rules (e.g., radio's special stage 4 handling)
    2. General rules (e.g., gripper correction policies)
    
    Only the first matching rule is applied per call.
    
    Args:
        task_id: Current task ID
        stage: Current stage
        state: Current state (23-dim)
        actions: Predicted actions (horizon, 23)
    
    Returns:
        (actions, stage) - corrected or original if no rule applies
    """
    # Try task-specific rules first
    task_specific_rules = TASK_SPECIFIC_RULES.get(task_id, [])
    for rule in task_specific_rules:
        result = rule(task_id, stage, state, actions)
        if result is not None:
            corrected_actions, corrected_stage = result
            return corrected_actions, corrected_stage
    
    # Try general rules (apply to all tasks)
    for rule in GENERAL_RULES:
        result = rule(task_id, stage, state, actions)
        if result is not None:
            corrected_actions, corrected_stage = result
            return corrected_actions, corrected_stage
    
    # No rules matched, return original
    return actions, stage


# ============================================================================
# GRIPPER VARIATION CHECK
# ============================================================================

def check_gripper_variation(
    actions: np.ndarray,
    num_actions_to_check: int
) -> Tuple[bool, float, float]:
    """
    Check if gripper actions have high variation.
    
    High variation in gripper positions suggests complex manipulation that
    should not be compressed via interpolation.
    
    Args:
        actions: Predicted actions (horizon, 23)
        num_actions_to_check: Number of actions to check (e.g., actions_to_execute)
    
    Returns:
        (has_high_variation, left_variation, right_variation)
    """
    # Check only the actions we plan to execute
    actions_to_check = actions[:num_actions_to_check]
    
    # Extract gripper actions
    left_gripper_actions = actions_to_check[:, LEFT_GRIPPER_IDX]
    right_gripper_actions = actions_to_check[:, RIGHT_GRIPPER_IDX]
    
    left_variation = float(np.max(left_gripper_actions) - np.min(left_gripper_actions))
    right_variation = float(np.max(right_gripper_actions) - np.min(right_gripper_actions))
    
    has_high_variation = (left_variation > GRIPPER_VARIATION_THRESHOLD or 
                         right_variation > GRIPPER_VARIATION_THRESHOLD)
    
    return has_high_variation, left_variation, right_variation

	"""
	Evaluation tricks for improving success rates on BEHAVIOR-1K tasks.

	This module contains:
	1. Correction rules: Task-specific problem resolution (stage reset, gripper recovery)
	2. Gripper variation check: Disable compression when grippers have large movement

	Enable via --apply-eval-tricks flag (default: False)
	"""

	import logging
	import numpy as np
	from typing import Tuple, Optional

	logger = logging.getLogger(__name__)


	# ============================================================================
	# GLOBAL CONSTANTS
	# ============================================================================

	# Gripper thresholds (normalized [-1, 1] range)
	# Grippers are normalized from raw [0, 0.1] to [-1, 1]:
	# -1.0 = fully closed (raw 0.0)
	# -0.9 = mostly closed (raw 0.005)
	# 0.0 = half open (raw 0.05)
	# 0.9 = mostly open (raw 0.095)
	# 1.0 = fully open (raw 0.1)
	OPEN_THRESHOLD = 0.90
	CLOSED_THRESHOLD = -0.98 # very strict threshold to avoid false positives

	# Gripper variation threshold for compression check
	GRIPPER_VARIATION_THRESHOLD = 0.2

	# State indices
	LEFT_GRIPPER_IDX = 14
	RIGHT_GRIPPER_IDX = 22
	LEFT_ARM_START_IDX = 7 # Left arm: state[7:14]
	RIGHT_ARM_START_IDX = 15 # Right arm: state[15:22]
	LEFT_ARM_END_IDX = 14
	RIGHT_ARM_END_IDX = 22

	# Base positions for arms (7 joints each)
	BASE_POSITION_LEFT_ARM = np.array([-0.2, 0.0, 0.1, -1.4, 0.1, 0.7, 0.0])
	BASE_POSITION_RIGHT_ARM = np.array([-0.2, 0.0, 0.1, -1.4, 0.1, 0.7, 0.0])


	# ============================================================================
	# GRIPPER CORRECTION POLICY
	# ============================================================================

	# Tasks where grippers should ALWAYS be open (any closure is an error)
	ALWAYS_OPEN_LEFT_GRIPPER_TASKS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 42, 43, 44, 45, 47, 48}
	ALWAYS_OPEN_RIGHT_GRIPPER_TASKS = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 34, 35, 36, 37, 42, 43, 44, 47, 48, 49}

	# Tasks where closure is not allowed BEFORE specified stage
	# Format: {task_id: {'left': min_stage, 'right': min_stage}}
	MIN_STAGE_FOR_CLOSURE = {
	0: {'left': 2, 'right': 2}, # Radio: allow closing from stage 2+
	30: {'right': 6}, # setting_the_fire
	31: {'right': 8}, # clean_boxing_gloves
	32: {'right': 5}, # wash_a_baseball_cap
	33: {'right': 11}, # wash_dog_toys
	40: {'right': 4}, # make_microwave_popcorn
	41: {'left': 14, 'right': 14}, # cook_cabbage
	45: {'right': 10}, # cook_hot_dogs
	46: {'left': 8, 'right': 8}, # cook_bacon
	49: {'left': 14}, # make_pizza
	}

	# Tasks where right gripper is ALWAYS allowed to close (exceptions)
	RIGHT_GRIPPER_ALWAYS_ALLOWED = {38, 39} # spraying_for_bugs, spraying_fruit_trees


	# ============================================================================
	# CORRECTION RULES - ADD NEW RULES HERE
	# ============================================================================

	def task0_stage4_reset_to_stage2(
	task_id: int,
	stage: int,
	state: np.ndarray,
	actions: np.ndarray
	) -> Optional[Tuple[np.ndarray, int]]:
	# Probably not very important for the final solution
	"""
	Task 0 (Radio), Stage 4: Reset to stage 2 with arm/gripper recovery.

	PRIMARY: task_id=0 AND stage=4 → reset to stage 2 (always)

	SECONDARY (conditional):
	- If both grippers are fully open or fully closed:
	→ Open closed grippers

	Args:
	task_id: Current task ID
	stage: Current stage
	state: Current state (23-dim)
	actions: Predicted actions (horizon, 23)

	Returns:
	(corrected_actions, corrected_stage) if rule applies, None otherwise
	"""
	# Check primary condition
	if task_id != 0 or stage < 2:
	return None

	# PRIMARY: Always reset stage
	if stage == 4:
	logger.info(f"🔧 CORRECTION RULE: Task 0 (Radio), Stage 4 → Stage 2")
	corrected_stage = 2
	else:
	corrected_stage = stage

	# SECONDARY: Check gripper states
	left_gripper = state[LEFT_GRIPPER_IDX]
	right_gripper = state[RIGHT_GRIPPER_IDX]

	left_closed = left_gripper < CLOSED_THRESHOLD
	right_closed = right_gripper < CLOSED_THRESHOLD
	left_open = left_gripper > OPEN_THRESHOLD
	right_open = right_gripper > OPEN_THRESHOLD

	# Check if any gripper is in middle position (not fully open or closed)
	left_middle = not (left_open or left_closed)
	right_middle = not (right_open or right_closed)

	change_action = False

	# Copy current state to all actions as baseline
	corrected_actions = np.tile(state, (actions.shape[0], 1))

	if left_closed and not right_middle:
	corrected_actions[:, LEFT_GRIPPER_IDX] = 1.0
	change_action = True
	logger.info(f" - Opening left gripper → 1.0")
	if right_closed and not left_middle:
	corrected_actions[:, RIGHT_GRIPPER_IDX] = 1.0
	change_action = True
	logger.info(f" - Opening right gripper → 1.0")

	if not change_action:
	if stage == corrected_stage:
	return None

	corrected_actions = actions

	return corrected_actions, corrected_stage


	def general_gripper_correction(
	task_id: int,
	stage: int,
	state: np.ndarray,
	actions: np.ndarray
	) -> Optional[Tuple[np.ndarray, int]]:
	"""
	General gripper correction rule for all tasks.

	Applies task-specific policies:
	1. Always-open tasks: Open gripper if closed (any stage)
	2. Stage-restricted tasks: Open gripper if closed before allowed stage
	3. Exceptions: Skip tasks where gripper is intentionally closed (e.g., spray bottles)

	Works in conjunction with task-specific rules:
	- Task 0: This rule handles stages ≤1, task-specific rule handles stage 4

	Args:
	task_id: Current task ID
	stage: Current stage
	state: np.ndarray,
	actions: Predicted actions (horizon, 23)

	Returns:
	(corrected_actions, corrected_stage) if rule applies, None otherwise
	"""
	# Get gripper states
	left_gripper = state[LEFT_GRIPPER_IDX]
	right_gripper = state[RIGHT_GRIPPER_IDX]

	left_closed = left_gripper < CLOSED_THRESHOLD
	right_closed = right_gripper < CLOSED_THRESHOLD

	# Determine if correction is needed
	left_needs_opening = False
	right_needs_opening = False

	# Check left gripper
	if left_closed:
	# Always open if task requires it
	if task_id in ALWAYS_OPEN_LEFT_GRIPPER_TASKS:
	left_needs_opening = True
	# Check stage restriction
	elif task_id in MIN_STAGE_FOR_CLOSURE:
	min_stage_left = MIN_STAGE_FOR_CLOSURE[task_id].get('left')
	if min_stage_left is not None and stage < min_stage_left:
	left_needs_opening = True

	# Check right gripper
	if right_closed:
	# Exception: Always allow for spray bottle tasks
	if task_id in RIGHT_GRIPPER_ALWAYS_ALLOWED:
	pass # Never correct these tasks
	# Always open if task requires it
	elif task_id in ALWAYS_OPEN_RIGHT_GRIPPER_TASKS:
	right_needs_opening = True
	# Check stage restriction
	elif task_id in MIN_STAGE_FOR_CLOSURE:
	min_stage_right = MIN_STAGE_FOR_CLOSURE[task_id].get('right')
	if min_stage_right is not None and stage < min_stage_right:
	right_needs_opening = True

	# Apply correction if needed
	if left_needs_opening or right_needs_opening:
	logger.info(f"🔧 GENERAL GRIPPER CORRECTION: Task {task_id}, Stage {stage}")

	# Copy current state to all actions
	corrected_actions = np.tile(state, (actions.shape[0], 1))

	if left_needs_opening:
	logger.info(f" Left gripper closed ({left_gripper:.4f}) → Opening to 1.0")
	corrected_actions[:, LEFT_GRIPPER_IDX] = 1.0

	if right_needs_opening:
	logger.info(f" Right gripper closed ({right_gripper:.4f}) → Opening to 1.0")
	corrected_actions[:, RIGHT_GRIPPER_IDX] = 1.0

	# Keep stage unchanged
	return corrected_actions, stage

	return None # No correction needed


	# ============================================================================
	# RULE REGISTRY - ORGANIZE BY TASK_ID
	# ============================================================================

	# Task-specific rules (these run BEFORE the general rule)
	# Dictionary mapping task_id -> list of rule functions
	TASK_SPECIFIC_RULES = {
	0: [task0_stage4_reset_to_stage2], # Radio has special stage 4 handling
	# Add more task-specific rules here as needed
	}

	# General rules that apply to ALL tasks (run after task-specific rules)
	GENERAL_RULES = [
	general_gripper_correction, # Applies to all tasks with gripper policies
	]


	# ============================================================================
	# RULE APPLICATION
	# ============================================================================

	def apply_correction_rules(
	task_id: int,
	stage: int,
	state: np.ndarray,
	actions: np.ndarray
	) -> Tuple[np.ndarray, int]:
	"""
	Apply correction rules for the given task_id.

	Rules are applied in order until one matches:
	1. Task-specific rules (e.g., radio's special stage 4 handling)
	2. General rules (e.g., gripper correction policies)

	Only the first matching rule is applied per call.

	Args:
	task_id: Current task ID
	stage: Current stage
	state: Current state (23-dim)
	actions: Predicted actions (horizon, 23)

	Returns:
	(actions, stage) - corrected or original if no rule applies
	"""
	# Try task-specific rules first
	task_specific_rules = TASK_SPECIFIC_RULES.get(task_id, [])
	for rule in task_specific_rules:
	result = rule(task_id, stage, state, actions)
	if result is not None:
	corrected_actions, corrected_stage = result
	return corrected_actions, corrected_stage

	# Try general rules (apply to all tasks)
	for rule in GENERAL_RULES:
	result = rule(task_id, stage, state, actions)
	if result is not None:
	corrected_actions, corrected_stage = result
	return corrected_actions, corrected_stage

	# No rules matched, return original
	return actions, stage


	# ============================================================================
	# GRIPPER VARIATION CHECK
	# ============================================================================

	def check_gripper_variation(
	actions: np.ndarray,
	num_actions_to_check: int
	) -> Tuple[bool, float, float]:
	"""
	Check if gripper actions have high variation.

	High variation in gripper positions suggests complex manipulation that
	should not be compressed via interpolation.

	Args:
	actions: Predicted actions (horizon, 23)
	num_actions_to_check: Number of actions to check (e.g., actions_to_execute)

	Returns:
	(has_high_variation, left_variation, right_variation)
	"""
	# Check only the actions we plan to execute
	actions_to_check = actions[:num_actions_to_check]

	# Extract gripper actions
	left_gripper_actions = actions_to_check[:, LEFT_GRIPPER_IDX]
	right_gripper_actions = actions_to_check[:, RIGHT_GRIPPER_IDX]

	left_variation = float(np.max(left_gripper_actions) - np.min(left_gripper_actions))
	right_variation = float(np.max(right_gripper_actions) - np.min(right_gripper_actions))

	has_high_variation = (left_variation > GRIPPER_VARIATION_THRESHOLD or
	right_variation > GRIPPER_VARIATION_THRESHOLD)

	return has_high_variation, left_variation, right_variation