Project: 006_MAP_Multi-Human-Value_Alignment_Palette | Function: optimize_lambda (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Core MAP optimization algorithm that solves the dual problem to find optimal lambda values for multi-human-value alignment using primal-dual optimization

**KEY TERMS:** {
  "_dual_objective": "dual objective function g(λ)",
  "lambda_vals": "λ (lambda parameters)",
  "self.c": "c (value palette constraints)",
  "self.rewards": "r(x,y) (reward vectors)",
  "Adam": "gradient-based optimization"
}

**CONTEXT CODE:**                 self.value_list = value_list.split(',')
        elif isinstance(value_list, list):
            self.value_list = value_list
        else:
            raise TypeError(f"value_list must be str or list, got {type(value_list)}")
        
        logger.info(f"Aligning values: {self.value_list}")
    
    def _process_rewards_tensor(self, rewards: Union[torch.Tensor, List[List[float]]]) -> None:
        """Process rewards provided as a tensor or list of lists.
        
        Args:
            rewards (Union[torch.Tensor, List[List[float]]]): Rewards tensor or list
        """
        if isinstance(rewards, list):
            rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device)
        
        if rewards.dim() == 1:
            # Single value, make 2D
            self.rewards = rewards.unsqueeze(0).to(self.device)
        else:
            self.rewards = rewards.to(self.device)
    
    def _process_rewards_data(self, rewards_data: Union[str, List[Dict[str, Any]]]) -> None:
        """Process rewards from a JSON file or list of dictionaries.
        
        Args:
            rewards_data (Union[str, List[Dict[str, Any]]]): JSON file path or data
        """
        # Load data if it's a file path
        if isinstance(rewards_data, str):
            with open(rewards_data, 'r') as file:
                data = json.load(file)
        else:
            data = rewards_data
        
        # Extract rewards for each value
        rewards_list = []
        
        for value in self.value_list:
            try:
                value_rewards = [entry[value] for entry in data]
                tensor_rewards = torch.tensor(value_rewards, dtype=torch.float32, device=self.device)
                rewards_list.append(tensor_rewards)
                logger.debug(f"Average {value} reward: {torch.mean(tensor_rewards).item():.4f}")
            except KeyError:
                raise ValueError(f"Value '{value}' not found in rewards data")
        
        self.rewards = torch.stack(rewards_list, dim=0)
    
    def optimize_lambda(
        self,
        lambda_init: Optional[List[float]] = None,
        optimize_indices: Optional[List[int]] = None,
        learning_rate: float = 0.1,
        max_steps: int = 150,
        verbose: bool = False,
        save_results: bool = False,
        save_path: Optional[str] = None
    ) -> Tuple[List[float], bool]:
        """Optimize lambda values for the given palette and rewards.

        Args:
            lambda_init (Optional[List[float]]): Initial lambda values
            optimize_indices (Optional[List[int]]): Indices of lambda values to optimize
            learning_rate (float): Learning rate for optimization
            max_steps (int): Maximum number of optimization steps
            verbose (bool): Whether to print detailed information
            save_results (bool): Whether to save results to file
            save_path (Optional[str]): Path to save results to

        Returns:
            Tuple[List[float], bool]: Optimized lambda values and success indicator
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def _dual_objective(self, lambda_vals: torch.Tensor) -> torch.Tensor:
        """Calculate the dual objective function for optimization.
        
        Args:
            lambda_vals (torch.Tensor): Lambda values
            
        Returns:
            torch.Tensor: Dual objective value
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    def sequential_optimize_lambda(
        self,
        lambda_init: Optional[List[float]] = None,
        learning_rate: float = 0.1,
        max_steps: int = 150,
        verbose: bool = False,
        rounds: int = 1
    ) -> List[float]:
        """Sequentially optimize lambda for each human value.

        Args:
            lambda_init (Optional[List[float]]): Initial lambda values
            learning_rate (float): Learning rate for optimization
            max_steps (int): Maximum number of steps for each optimization
            verbose (bool): Whether to print detailed information
            rounds (int): Number of rounds of sequential optimization

        Returns:
            List[float]: Optimized lambda values
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def _save_results(
        self, 
        lambda_vals: List[float], 
        success: bool,
        save_path: Optional[str] = None
    ) -> None:
        """Save optimization results to a file.
        
        Args:
            lambda_vals (List[float]): Optimized lambda values
            success (bool): Whether optimization was successful
            save_path (Optional[str]): Path to save results to
        """
        c_str = ','.join(f'{v:.3f}' for v in self.c.tolist())
        lambda_str = ','.join(f'{v:.3f}' for v in lambda_vals) if success else 'NaN'
        values_str = ','.join(self.value_list)
        
        result_line = f"values: {values_str}, target_palette: {c_str}, optimized_lambda: {lambda_str}\n"
        
        # Default save path
        if save_path is None:
            save_path = "alignmap_results.txt"
        
        # Create directory if it doesn't exist
        os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
        
        # Append results to file
        with open(save_path, 'a') as file:
            file.write(result_line)
        
        logger.info(f"Results saved to {save_path}")
        
    def _save_results_to_text(
        self, 
        lambda_vals: List[float], 
        success: bool,
        save_prefix: str = 'results/alignValues'
    ) -> None:
        """Save optimization results to a text file (legacy format).
        
        This method maintains backward compatibility with the original alignValues.py
        result saving format.


Complete the function implementation replacing "raise NotImplementedError".