dso_experiments.py

"""
Deep Symbolic Regression Experiments using DSO Library

This script runs symbolic regression experiments using an RNN-based approach
with reinforcement learning. Three training policies are implemented:
1. Risk-Seeking Policy Gradient
2. Vanilla Policy Gradient
3. Priority Queue Training

Uses training data and splits from ABGGeneric.ipynb for the ABG propagation model.

Requirements:
    pip install deep-symbolic-optimization pandas numpy scikit-learn openpyxl matplotlib

Usage:
    python dso_experiments.py --policy all  # Run all three policies
    python dso_experiments.py --policy risk_seeking  # Run only risk-seeking
    python dso_experiments.py --policy vanilla  # Run only vanilla PG
    python dso_experiments.py --policy priority_queue  # Run only PQT
"""

import os
import re
import json
import argparse
import warnings
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Tuple, Optional, Any

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, mean_absolute_percentage_error

warnings.filterwarnings('ignore')

# ============================================================================
# Data Loading and Preprocessing (matching ABGGeneric.ipynb)
# ============================================================================

def extract_interval(interval_value: Any) -> Tuple[float, float]:
    """Extract min/max values from interval notation like '[x,y]' or [x, y]."""
    interval_str = str(interval_value)

    if isinstance(interval_str, str):
        match = re.match(r'\[([0-9.-]+),\s*([0-9.-]+)\]', interval_str)
        if match:
            return float(match.group(1)), float(match.group(2))
        else:
            raise ValueError(f"Invalid interval format: {interval_str}")
    elif isinstance(interval_value, list) and len(interval_value) == 2:
        return float(interval_value[0]), float(interval_value[1])
    else:
        raise ValueError(f"Invalid interval format: {interval_value}")


def load_abg_data(excel_path: str = 'Propagation_models.xlsx') -> pd.DataFrame:
    """
    Load ABG data from Excel file with normalization as in ABGGeneric.ipynb.

    Args:
        excel_path: Path to the Excel file with propagation models

    Returns:
        Normalized DataFrame with ABG data
    """
    # Read the data from the Excel file
    abg_data = pd.read_excel(excel_path, sheet_name='ABG', header=0)

    # First row contains normalization factors
    first_row = abg_data.iloc[0]

    # Drop the first row after using it for normalization
    abg_data = abg_data.drop(index=0)

    # Clean column names by stripping whitespace
    abg_data.columns = abg_data.columns.str.strip()

    # Normalize the columns based on the first row
    feature_cols = ['alpha', 'gamma', 'beta', 'f', 'd', 'X']
    normalization_info = {}

    for col in feature_cols:
        min_val, max_val = extract_interval(first_row[col])
        normalization_info[col] = {'min': min_val, 'max': max_val}
        abg_data[col] = (abg_data[col] - min_val) / (max_val - min_val)

    # Convert specified columns to floats
    for col in feature_cols:
        abg_data[col] = pd.to_numeric(abg_data[col], errors='coerce').astype(float)

    # Check for any NaN values after conversion
    if abg_data.isnull().any().any():
        print("Warning: NaN values detected in the dataset. Dropping rows with NaN.")
        abg_data = abg_data.dropna()

    return abg_data, normalization_info


def load_ci_data(excel_path: str = 'Propagation_models.xlsx') -> pd.DataFrame:
    """
    Load CI (Close-In) data from Excel file with normalization.

    Args:
        excel_path: Path to the Excel file with propagation models

    Returns:
        Normalized DataFrame with CI data and normalization info
    """
    # Read the data from the Excel file
    ci_data = pd.read_excel(excel_path, sheet_name='CI', header=0)

    # First row contains normalization factors
    first_row = ci_data.iloc[0]

    # Drop the first row after using it for normalization
    ci_data = ci_data.drop(index=0)

    # Clean column names by stripping whitespace
    ci_data.columns = ci_data.columns.str.strip()

    # Normalize the columns based on the first row
    feature_cols = ['f', 'd', 'X', 'n']
    normalization_info = {}

    for col in feature_cols:
        min_val, max_val = extract_interval(first_row[col])
        normalization_info[col] = {'min': min_val, 'max': max_val}
        ci_data[col] = (ci_data[col] - min_val) / (max_val - min_val)

    # Convert specified columns to floats
    for col in feature_cols:
        ci_data[col] = pd.to_numeric(ci_data[col], errors='coerce').astype(float)

    # Check for any NaN values after conversion
    if ci_data.isnull().any().any():
        print("Warning: NaN values detected in the dataset. Dropping rows with NaN.")
        ci_data = ci_data.dropna()

    return ci_data, normalization_info


def load_indoor_data(excel_path: str = 'csv_experiments.xlsx') -> Tuple[pd.DataFrame, Dict]:
    """
    Load Indoor data from Excel file with normalization.

    Args:
        excel_path: Path to the Excel file with propagation models

    Returns:
        Normalized DataFrame with Indoor data and normalization info
    """
    # Read the data from the Excel file
    indoor_data = pd.read_excel(excel_path, sheet_name='Indoor', header=0)

    # Clean column names by stripping whitespace
    indoor_data.columns = indoor_data.columns.str.strip()

    # Normalize the columns based on computed min/max (no normalization row in this dataset)
    feature_cols = ['nw', 'nf', 'd', 'f']
    normalization_info = {}

    for col in feature_cols:
        min_val, max_val = indoor_data[col].min(), indoor_data[col].max()
        normalization_info[col] = {'min': float(min_val), 'max': float(max_val)}
        indoor_data[col] = (indoor_data[col] - min_val) / (max_val - min_val)

    # Convert specified columns to floats
    for col in feature_cols:
        indoor_data[col] = pd.to_numeric(indoor_data[col], errors='coerce').astype(float)

    # Check for any NaN values after conversion
    if indoor_data.isnull().any().any():
        print("Warning: NaN values detected in the dataset. Dropping rows with NaN.")
        indoor_data = indoor_data.dropna()

    return indoor_data, normalization_info


def load_outdoor_data(excel_path: str = 'csv_experiments.xlsx') -> Tuple[pd.DataFrame, Dict]:
    """
    Load Outdoor data from Excel file with normalization.

    Args:
        excel_path: Path to the Excel file with propagation models

    Returns:
        Normalized DataFrame with Outdoor data and normalization info
    """
    # Read the data from the Excel file
    outdoor_data = pd.read_excel(excel_path, sheet_name='Outdoor', header=0)

    # Clean column names by stripping whitespace
    outdoor_data.columns = outdoor_data.columns.str.strip()

    # Normalize the columns based on computed min/max (no normalization row in this dataset)
    feature_cols = ['hed', 'd', 'f']
    normalization_info = {}

    for col in feature_cols:
        min_val, max_val = outdoor_data[col].min(), outdoor_data[col].max()
        normalization_info[col] = {'min': float(min_val), 'max': float(max_val)}
        outdoor_data[col] = (outdoor_data[col] - min_val) / (max_val - min_val)

    # Convert specified columns to floats
    for col in feature_cols:
        outdoor_data[col] = pd.to_numeric(outdoor_data[col], errors='coerce').astype(float)

    # Check for any NaN values after conversion
    if outdoor_data.isnull().any().any():
        print("Warning: NaN values detected in the dataset. Dropping rows with NaN.")
        outdoor_data = outdoor_data.dropna()

    return outdoor_data, normalization_info


# Dataset configuration for feature columns and target columns
DATASET_CONFIG = {
    'abg': {'feature_cols': ['alpha', 'gamma', 'beta', 'f', 'd', 'X'], 'target_col': 'PL_ABG'},
    'ci': {'feature_cols': ['f', 'd', 'X', 'n'], 'target_col': 'PL_CI'},
    'indoor': {'feature_cols': ['nw', 'nf', 'd', 'f'], 'target_col': 'PL'},
    'outdoor': {'feature_cols': ['hed', 'd', 'f'], 'target_col': 'PL'}
}


def prepare_train_test_split(
    data: pd.DataFrame,
    dataset_type: str = 'abg',
    test_size: float = 0.2,
    val_size: float = 0.2,
    random_state: int = 42
) -> Dict[str, np.ndarray]:
    """
    Prepare train/validation/test splits for ABG or CI data.

    Args:
        data: DataFrame with propagation model data
        dataset_type: Type of dataset ('abg' or 'ci')
        test_size: Fraction of data for test set
        val_size: Fraction of train data for validation
        random_state: Random seed for reproducibility

    Returns:
        Dictionary with train/val/test inputs and labels
    """
    # Get feature and target columns based on dataset type
    config = DATASET_CONFIG[dataset_type]
    feature_cols = config['feature_cols']
    target_col = config['target_col']

    # Extract inputs and labels
    inputs = data[feature_cols].values
    labels = data[target_col].values

    # Split into 80% train + 20% test
    train_val_input, test_input, train_val_label, test_label = train_test_split(
        inputs, labels, test_size=test_size, random_state=random_state
    )

    # Split the remaining 80% into 80% train + 20% validation
    train_input, val_input, train_label, val_label = train_test_split(
        train_val_input, train_val_label, test_size=val_size, random_state=random_state
    )

    return {
        'train_input': train_input,
        'train_label': train_label,
        'val_input': val_input,
        'val_label': val_label,
        'test_input': test_input,
        'test_label': test_label,
        'feature_names': feature_cols,
        'target_name': target_col
    }


def save_data_for_dso(data_splits: Dict, output_dir: str = 'dso_data') -> Dict[str, str]:
    """
    Save data splits as CSV files for DSO library.

    Args:
        data_splits: Dictionary with train/val/test data
        output_dir: Directory to save CSV files

    Returns:
        Dictionary with paths to saved files
    """
    os.makedirs(output_dir, exist_ok=True)

    feature_names = data_splits['feature_names']
    target_name = data_splits.get('target_name', 'PL_ABG')
    paths = {}

    for split_name in ['train', 'val', 'test']:
        input_key = f'{split_name}_input'
        label_key = f'{split_name}_label'

        # Combine inputs and labels
        df = pd.DataFrame(data_splits[input_key], columns=feature_names)
        df[target_name] = data_splits[label_key]

        # Save to CSV (no header - DSO expects header=None format)
        csv_path = os.path.join(output_dir, f'{split_name}_data.csv')
        df.to_csv(csv_path, index=False, header=False)
        paths[split_name] = csv_path
        print(f"Saved {split_name} data to {csv_path} ({len(df)} samples)")

    # Also create a combined training dataset (train + val) for DSO
    train_df = pd.DataFrame(data_splits['train_input'], columns=feature_names)
    train_df[target_name] = data_splits['train_label']

    val_df = pd.DataFrame(data_splits['val_input'], columns=feature_names)
    val_df[target_name] = data_splits['val_label']

    combined_df = pd.concat([train_df, val_df], ignore_index=True)
    combined_path = os.path.join(output_dir, 'train_val_data.csv')
    combined_df.to_csv(combined_path, index=False, header=False)
    paths['train_val'] = combined_path
    print(f"Saved combined train+val data to {combined_path} ({len(combined_df)} samples)")

    return paths


# ============================================================================
# DSO Configuration Generators
# ============================================================================

def create_base_config(
    dataset_path: str,
    experiment_name: str,
    n_samples: int = 2000000,
    batch_size: int = 200,
    dataset_type: str = 'abg',
) -> Dict:
    """
    Create base DSO configuration common to all policies.

    Args:
        dataset_path: Path to training CSV file
        experiment_name: Name for this experiment
        n_samples: Total number of samples to generate during training
        batch_size: Batch size for training (default: 200)
        dataset_type: Type of dataset ('abg' or 'ci')

    Returns:
        Base configuration dictionary
    """
    # Repeat constraint tokens depend on dataset type
    # ABG features: ['alpha', 'gamma', 'beta', 'f', 'd', 'X'] -> f=x3, d=x4
    # CI features: ['f', 'd', 'X', 'n'] -> f=x0, d=x1
    # Indoor features: ['nw', 'nf', 'd', 'f'] -> d=x2, f=x3
    # Outdoor features: ['hed', 'd', 'f'] -> d=x1, f=x2
    repeat_tokens = {
        'abg': ["x3", "x4"],     # f (index 3) and d (index 4) in ABG feature order
        'ci': ["x0", "x1"],      # f (index 0) and d (index 1) in CI feature order
        'indoor': ["x2", "x3"],  # d (index 2) and f (index 3) in Indoor feature order
        'outdoor': ["x1", "x2"]  # d (index 1) and f (index 2) in Outdoor feature order
    }

    return {
        "experiment": {
            "logdir": "./logs",
            "exp_name": experiment_name,
            "seed": 42
        },
        "task": {
            "task_type": "regression",
            "dataset": dataset_path,
            "function_set": [
                "add", "sub", "mul", "div",
                "sin", "cos", "tan",
                "exp", "log", "sqrt",
                "n2", "n3", "neg", "inv",
                "const"
            ],
            "metric": "inv_nmse",
            "metric_params": [1.0],
            "threshold": 1e-12,
            "protected": True
        },
        "training": {
            "n_samples": n_samples,
            "batch_size": batch_size,
            "alpha": 0.5,
            "verbose": True
        },
        "policy": {
            "policy_type": "rnn",
            "max_length": 64,
            "cell": "lstm",
            "num_layers": 2,
            "num_units": 64,
            "initializer": "zeros"
        },
        "prior": {
            # Constraint 1: Expression length between 4 and 40 tokens
            "length": {
                "min_": 4,
                "max_": 40,
                "on": True
            },
            # Constraint 2: Repeat constraint for f and d variables
            "repeat": {
                "tokens": repeat_tokens.get(dataset_type, ["x3", "x4"]),
                "min_": None,
                "max_": 3,
                "on": True
            },
            # Constraint 3: Inverse of unary operations should not be children
            # (e.g., log(exp(x)) or exp(log(x)) are discouraged)
            "inverse": {
                "on": True
            },
            # Constraint 4: Trigonometric operators should not have other trig operators as descendants
            "trig": {
                "on": True
            },
            # Constraint 5: Restrict operators to prevent all children from being constants
            "no_inputs": {
                "on": True
            },
            # Constant token constraint
            "const": {
                "on": True
            }
        }
    }


def create_risk_seeking_config(
    dataset_path: str,
    experiment_name: str = "abg_risk_seeking",
    epsilon: float = 0.05,
    learning_rate: float = 0.002,
    entropy_weight: float = 0.008,
    batch_size: int = 200,
    dataset_type: str = 'abg',
    **kwargs
) -> Dict:
    """
    Create configuration for Risk-Seeking Policy Gradient.

    Risk-seeking PG only updates parameters on top-performing designs
    using quantile-based thresholding. This focuses learning on the
    most promising symbolic expressions.

    Args:
        dataset_path: Path to training CSV file
        experiment_name: Name for this experiment
        epsilon: Risk factor (fraction of top samples to use, e.g., 0.05 = top 5%)
        learning_rate: Learning rate for policy optimizer (default: 0.002)
        entropy_weight: Entropy regularization weight (default: 0.008)
        dataset_type: Type of dataset ('abg' or 'ci')
        **kwargs: Additional arguments passed to base config

    Returns:
        Configuration dictionary for risk-seeking policy gradient
    """
    config = create_base_config(dataset_path, experiment_name, batch_size=batch_size, dataset_type=dataset_type, **kwargs)

    # Risk-seeking specific settings
    config["training"].update({
        "epsilon": epsilon,  # Risk factor: only use top epsilon fraction
        "baseline": "R_e",   # Risk-seeking baseline (quantile threshold)
    })

    config["policy_optimizer"] = {
        "policy_optimizer_type": "pg",  # Policy gradient
        "learning_rate": learning_rate,
        "entropy_weight": entropy_weight,
        "entropy_gamma": 1.0
    }

    return config


def create_vanilla_pg_config(
    dataset_path: str,
    experiment_name: str = "abg_vanilla_pg",
    learning_rate: float = 0.0001,
    entropy_weight: float = 0.005,
    batch_size: int = 200,
    dataset_type: str = 'abg',
    **kwargs
) -> Dict:
    """
    Create configuration for Vanilla Policy Gradient.

    Vanilla PG uses an exponential weighted moving average baseline
    and learns from all generated samples, not just the best ones.

    Args:
        dataset_path: Path to training CSV file
        experiment_name: Name for this experiment
        learning_rate: Learning rate for policy optimizer (default: 0.002)
        entropy_weight: Entropy regularization weight (default: 0.008)
        dataset_type: Type of dataset ('abg' or 'ci')
        **kwargs: Additional arguments passed to base config

    Returns:
        Configuration dictionary for vanilla policy gradient
    """
    config = create_base_config(dataset_path, experiment_name, batch_size=batch_size, dataset_type=dataset_type, **kwargs)

    # Vanilla PG specific settings
    config["training"].update({
        "epsilon": 1.0,      # Use all samples (no risk-seeking)
        "baseline": "ewma_R",  # Exponential weighted moving average baseline
    })

    config["policy_optimizer"] = {
        "policy_optimizer_type": "pg",  # Policy gradient
        "learning_rate": learning_rate,
        "entropy_weight": entropy_weight,
        "entropy_gamma": 0.99
    }

    return config


def create_priority_queue_config(
    dataset_path: str,
    experiment_name: str = "abg_priority_queue",
    pq_k: int = 10,
    learning_rate: float = 0.002,
    entropy_weight: float = 0.005,
    batch_size: int = 200,
    dataset_type: str = 'abg',
    **kwargs
) -> Dict:
    """
    Create configuration for Priority Queue Training (PQT).

    PQT maintains a priority queue of the best expressions discovered
    and trains only on these top-K highest-reward designs.

    Args:
        dataset_path: Path to training CSV file
        experiment_name: Name for this experiment
        pq_k: Size of priority queue (number of top expressions to keep)
        learning_rate: Learning rate for policy optimizer (default: 0.002)
        entropy_weight: Entropy regularization weight (default: 0.008)
        dataset_type: Type of dataset ('abg' or 'ci')
        **kwargs: Additional arguments passed to base config

    Returns:
        Configuration dictionary for priority queue training
    """
    config = create_base_config(dataset_path, experiment_name, batch_size=batch_size, dataset_type=dataset_type, **kwargs)

    # Priority Queue specific settings - use ewma_R baseline for compatibility
    config["training"].update({
        "epsilon": 1.0,
        "baseline": "ewma_R",
    })

    config["policy_optimizer"] = {
        "policy_optimizer_type": "pqt",  # Priority queue training
        "learning_rate": learning_rate,
        "entropy_weight": entropy_weight,
        "pqt_k": pq_k,  # Keep top-k expressions
        "pqt_batch_size": pq_k,  # Batch size for PQ updates
    }

    return config


def save_config(config: Dict, output_path: str) -> str:
    """Save configuration to JSON file."""
    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
    with open(output_path, 'w') as f:
        json.dump(config, f, indent=2)
    print(f"Saved config to {output_path}")
    return output_path


# ============================================================================
# Experiment Runner
# ============================================================================

class DSOExperiment:
    """
    Class to run and manage DSO experiments with different policies.
    """

    POLICIES = {
        'risk_seeking': create_risk_seeking_config,
        'vanilla': create_vanilla_pg_config,
        'priority_queue': create_priority_queue_config
    }

    # Policy-specific default parameters for ABG dataset
    ABG_POLICY_PARAMS = {
        'risk_seeking': {'learning_rate': 0.002, 'entropy_weight': 0.008, 'batch_size': 200},
        'vanilla': {'learning_rate': 0.0001, 'entropy_weight': 0.005, 'batch_size': 200},
        'priority_queue': {'learning_rate': 0.002, 'entropy_weight': 0.005, 'batch_size': 200}
    }

    # Policy-specific default parameters for CI dataset
    CI_POLICY_PARAMS = {
        'risk_seeking': {'learning_rate': 0.001, 'entropy_weight': 0.008, 'batch_size': 200},
        'vanilla': {'learning_rate': 0.0005, 'entropy_weight': 0.008, 'batch_size': 200},
        'priority_queue': {'learning_rate': 0.002, 'entropy_weight': 0.005, 'batch_size': 200}
    }

    # Policy-specific default parameters for Indoor dataset
    INDOOR_POLICY_PARAMS = {
        'risk_seeking': {'learning_rate': 0.0005, 'entropy_weight': 0.03, 'batch_size': 300},
        'vanilla': {'learning_rate': 0.001, 'entropy_weight': 0.02, 'batch_size': 200},
        'priority_queue': {'learning_rate': 0.001, 'entropy_weight': 0.01, 'batch_size': 200}
    }

    # Policy-specific default parameters for Outdoor dataset
    OUTDOOR_POLICY_PARAMS = {
        'risk_seeking': {'learning_rate': 0.0005, 'entropy_weight': 0.01, 'batch_size': 200},
        'vanilla': {'learning_rate': 0.0001, 'entropy_weight': 0.01, 'batch_size': 200},
        'priority_queue': {'learning_rate': 0.0005, 'entropy_weight': 0.01, 'batch_size': 200}
    }

    # Mapping of dataset types to their policy parameters
    DATASET_POLICY_PARAMS = {
        'abg': ABG_POLICY_PARAMS,
        'ci': CI_POLICY_PARAMS,
        'indoor': INDOOR_POLICY_PARAMS,
        'outdoor': OUTDOOR_POLICY_PARAMS
    }

    def __init__(
        self,
        data_path: str = 'Propagation_models.xlsx',
        output_dir: str = 'dso_experiments',
        n_samples: int = 50000,
        batch_size: int = 200,
        n_runs: int = 10,
        random_seed: int = 42,
        dataset_type: str = 'abg',
        policy_params: Optional[Dict[str, Dict]] = None
    ):
        """
        Initialize DSO experiment manager.

        Args:
            data_path: Path to Excel file with propagation model data
            output_dir: Directory for experiment outputs
            n_samples: Number of samples per training run
            batch_size: Batch size for training (default: 200)
            n_runs: Number of independent runs per policy (default: 10)
            random_seed: Base random seed
            dataset_type: Type of dataset ('abg' or 'ci')
            policy_params: Optional dict mapping policy names to their specific parameters
        """
        self.data_path = data_path
        self.output_dir = output_dir
        self.n_samples = n_samples
        self.batch_size = batch_size
        self.n_runs = n_runs
        self.random_seed = random_seed
        self.dataset_type = dataset_type

        # Select policy params based on dataset type and merge with any custom ones
        default_policy_params = self.DATASET_POLICY_PARAMS.get(dataset_type, self.ABG_POLICY_PARAMS)
        self.policy_params = {k: dict(v) for k, v in default_policy_params.items()}
        if policy_params:
            for policy, params in policy_params.items():
                if policy in self.policy_params:
                    self.policy_params[policy].update(params)
                else:
                    self.policy_params[policy] = params

        # Create output directories
        self.data_dir = os.path.join(output_dir, 'data')
        self.config_dir = os.path.join(output_dir, 'configs')
        self.results_dir = os.path.join(output_dir, 'results')
        self.logs_dir = os.path.join(output_dir, 'logs')

        for d in [self.data_dir, self.config_dir, self.results_dir, self.logs_dir]:
            os.makedirs(d, exist_ok=True)

        self.data_splits = None
        self.data_paths = None
        self.results = {}

    def prepare_data(self) -> None:
        """Load and prepare data for DSO experiments."""
        print("=" * 60)
        print(f"Loading and preparing {self.dataset_type.upper()} data...")
        print("=" * 60)

        # Load data based on dataset type
        if self.dataset_type == 'indoor':
            data, norm_info = load_indoor_data(self.data_path)
        elif self.dataset_type == 'outdoor':
            data, norm_info = load_outdoor_data(self.data_path)
        elif self.dataset_type == 'ci':
            data, norm_info = load_ci_data(self.data_path)
        else:
            data, norm_info = load_abg_data(self.data_path)
        print(f"Loaded {len(data)} samples from {self.data_path}")

        # Save normalization info
        norm_path = os.path.join(self.data_dir, 'normalization_info.json')
        with open(norm_path, 'w') as f:
            json.dump(norm_info, f, indent=2)

        # Create train/val/test splits
        self.data_splits = prepare_train_test_split(
            data,
            dataset_type=self.dataset_type,
            random_state=self.random_seed
        )
        print(f"Train: {len(self.data_splits['train_input'])} samples")
        print(f"Val: {len(self.data_splits['val_input'])} samples")
        print(f"Test: {len(self.data_splits['test_input'])} samples")

        # Save data as CSV for DSO
        self.data_paths = save_data_for_dso(self.data_splits, self.data_dir)

    def create_configs(self, policies: Optional[List[str]] = None) -> Dict[str, str]:
        """
        Create configuration files for specified policies.

        Args:
            policies: List of policy names to create configs for.
                     If None, creates configs for all policies.

        Returns:
            Dictionary mapping policy names to config file paths
        """
        if policies is None:
            policies = list(self.POLICIES.keys())

        config_paths = {}

        print("\n" + "=" * 60)
        print("Creating configuration files...")
        print("=" * 60)

        for policy_name in policies:
            if policy_name not in self.POLICIES:
                print(f"Warning: Unknown policy '{policy_name}', skipping...")
                continue

            config_func = self.POLICIES[policy_name]

            # Get policy-specific parameters
            policy_specific = self.policy_params.get(policy_name, {})

            config = config_func(
                dataset_path=os.path.abspath(self.data_paths['train_val']),
                experiment_name=f"{self.dataset_type}_{policy_name}",
                n_samples=self.n_samples,
                dataset_type=self.dataset_type,
                **policy_specific
            )

            # Update log directory
            config['experiment']['logdir'] = os.path.abspath(self.logs_dir)

            config_path = os.path.join(self.config_dir, f'{policy_name}_config.json')
            save_config(config, config_path)
            config_paths[policy_name] = config_path

            # Print policy-specific parameters
            print(f"  {policy_name}: lr={policy_specific.get('learning_rate', 'default')}, "
                  f"entropy={policy_specific.get('entropy_weight', 'default')}, "
                  f"batch={policy_specific.get('batch_size', 'default')}")

        return config_paths

    def run_single_experiment(
        self,
        policy_name: str,
        config_path: str,
        run_id: int = 0
    ) -> Dict:
        """
        Run a single DSO experiment.

        Args:
            policy_name: Name of the policy being used
            config_path: Path to configuration file
            run_id: Run identifier for multiple runs

        Returns:
            Dictionary with experiment results
        """
        print(f"\n--- Running {policy_name} (run {run_id + 1}/{self.n_runs}) ---")

        try:
            from dso import DeepSymbolicOptimizer

            # Create optimizer and train
            seed = self.random_seed + run_id * 100
            model = DeepSymbolicOptimizer(config_path)
            model.config['experiment']['seed'] = seed

            # Train the model
            result = model.train()

            # Extract best program - try multiple possible keys
            best_program = None
            best_reward = None

            # Try different keys for program extraction
            if result is not None:
                # Try 'program' key first, then alternatives
                best_program = result.get('program') or result.get('expression') or result.get('best_program')

                # Try different keys for reward
                best_reward = result.get('reward') or result.get('r') or result.get('best_reward')

            # If still no program, try to extract from Hall of Fame
            if best_program is None:
                print("  Warning: Could not extract program from result, checking Hall of Fame...")
                best_program, best_reward = self._extract_from_hall_of_fame(policy_name, run_id)

            if best_program is None:
                raise ValueError(f"Could not extract best program from DSO result. "
                               f"Result keys: {list(result.keys()) if result else 'None'}")

            # Evaluate on test set
            test_metrics = self._evaluate_program(
                best_program,
                self.data_splits['test_input'],
                self.data_splits['test_label']
            )

            return {
                'policy': policy_name,
                'run_id': run_id,
                'seed': seed,
                'best_expression': str(best_program),
                'best_reward': float(best_reward) if best_reward is not None else None,
                'test_metrics': test_metrics,
                'success': True
            }

        except ImportError:
            print("Warning: DSO library not installed. Running in simulation mode.")
            return self._simulate_experiment(policy_name, run_id)
        except Exception as e:
            print(f"Error during experiment: {e}")
            import traceback
            traceback.print_exc()
            return {
                'policy': policy_name,
                'run_id': run_id,
                'success': False,
                'error': str(e)
            }

    def _extract_from_hall_of_fame(
        self,
        policy_name: str,
        run_id: int
    ) -> Tuple[Optional[Any], Optional[float]]:
        """
        Extract best program from Hall of Fame CSV file as fallback.

        Args:
            policy_name: Name of the policy
            run_id: Run identifier

        Returns:
            Tuple of (best_program, best_reward) or (None, None) if not found
        """
        # Look for Hall of Fame files in logs directory
        import glob

        hof_patterns = [
            os.path.join(self.logs_dir, f"*{policy_name}*", "hof*.csv"),
            os.path.join(self.logs_dir, f"abg_{policy_name}*", "hof*.csv"),
            os.path.join(self.logs_dir, "**", "hof*.csv"),
        ]

        for pattern in hof_patterns:
            hof_files = glob.glob(pattern, recursive=True)
            if hof_files:
                # Get most recent file
                hof_file = max(hof_files, key=os.path.getmtime)
                try:
                    hof_df = pd.read_csv(hof_file)
                    if len(hof_df) > 0:
                        # Get best row (highest reward/r)
                        reward_col = 'r' if 'r' in hof_df.columns else 'reward'
                        if reward_col in hof_df.columns:
                            best_idx = hof_df[reward_col].idxmax()
                            best_row = hof_df.loc[best_idx]

                            # Get expression
                            expr_col = 'expression' if 'expression' in hof_df.columns else 'traversal'
                            expression = best_row.get(expr_col, str(best_row))
                            reward = best_row.get(reward_col, None)

                            print(f"  Extracted from Hall of Fame: {hof_file}")
                            return expression, reward
                except Exception as e:
                    print(f"  Warning: Could not parse Hall of Fame file {hof_file}: {e}")

        return None, None

    def _evaluate_program(
        self,
        program,
        X: np.ndarray,
        y: np.ndarray
    ) -> Dict[str, float]:
        """Evaluate a symbolic program on test data."""
        try:
            y_pred = program.execute(X)
            mape = mean_absolute_percentage_error(y, y_pred) * 100  # Convert to percentage
            return {
                'mse': float(mean_squared_error(y, y_pred)),
                'rmse': float(np.sqrt(mean_squared_error(y, y_pred))),
                'mae': float(mean_absolute_error(y, y_pred)),
                'mape': float(mape),
                'r2': float(r2_score(y, y_pred))
            }
        except Exception as e:
            return {'error': str(e)}

    def _simulate_experiment(self, policy_name: str, run_id: int) -> Dict:
        """Simulate experiment results when DSO is not installed."""
        # Simulated results for testing the experiment framework
        np.random.seed(self.random_seed + run_id)

        base_r2 = {'risk_seeking': 0.85, 'vanilla': 0.80, 'priority_queue': 0.82}
        r2 = base_r2.get(policy_name, 0.75) + np.random.normal(0, 0.05)
        mse = (1 - r2) * 100 + np.random.normal(0, 5)
        mape = (1 - r2) * 15 + np.random.normal(0, 2)

        return {
            'policy': policy_name,
            'run_id': run_id,
            'seed': self.random_seed + run_id * 100,
            'best_expression': f"alpha + beta*log(f) + gamma*d + X (simulated for {policy_name})",
            'best_reward': float(r2),
            'test_metrics': {
                'mse': float(max(0.1, mse)),
                'rmse': float(np.sqrt(max(0.1, mse))),
                'mae': float(max(0.1, np.sqrt(mse) * 0.8)),
                'mape': float(max(0.1, mape)),
                'r2': float(min(0.99, max(0.5, r2)))
            },
            'success': True,
            'simulated': True
        }

    def run_experiments(
        self,
        policies: Optional[List[str]] = None
    ) -> Dict[str, List[Dict]]:
        """
        Run experiments for all specified policies.

        Args:
            policies: List of policies to run. If None, runs all.

        Returns:
            Dictionary mapping policy names to lists of run results
        """
        if self.data_splits is None:
            self.prepare_data()

        config_paths = self.create_configs(policies)

        if policies is None:
            policies = list(config_paths.keys())

        print("\n" + "=" * 60)
        print("Running experiments...")
        print("=" * 60)

        for policy_name in policies:
            if policy_name not in config_paths:
                continue

            self.results[policy_name] = []

            for run_id in range(self.n_runs):
                result = self.run_single_experiment(
                    policy_name,
                    config_paths[policy_name],
                    run_id
                )
                self.results[policy_name].append(result)

                # Save intermediate results
                self._save_results()

        return self.results

    def _save_results(self) -> None:
        """Save current results to JSON file."""
        results_path = os.path.join(self.results_dir, 'experiment_results.json')
        with open(results_path, 'w') as f:
            json.dump(self.results, f, indent=2)
        # Also save per-run metrics CSV
        self.save_run_metrics_csv()

    def save_run_metrics_csv(self) -> None:
        """Save per-run metrics to CSV for easy averaging across runs."""
        rows = []
        for policy, runs in self.results.items():
            for run in runs:
                if run.get('success') and 'test_metrics' in run:
                    m = run['test_metrics']
                    rows.append({
                        'policy': policy,
                        'run_id': run['run_id'],
                        'seed': run.get('seed'),
                        'mae': m.get('mae'),
                        'mse': m.get('mse'),
                        'mape': m.get('mape'),
                        'r2': m.get('r2'),
                        'expression': run.get('best_expression', '')
                    })

        if rows:
            df = pd.DataFrame(rows)
            csv_path = os.path.join(self.results_dir, 'run_metrics.csv')
            df.to_csv(csv_path, index=False)
            print(f"Per-run metrics saved to {csv_path}")

    def summarize_results(self) -> pd.DataFrame:
        """
        Create summary DataFrame of experiment results.

        Returns:
            DataFrame with aggregated metrics per policy
        """
        if not self.results:
            print("No results to summarize. Run experiments first.")
            return pd.DataFrame()

        summary_data = []

        for policy_name, runs in self.results.items():
            successful_runs = [r for r in runs if r.get('success', False)]

            if not successful_runs:
                continue

            metrics = {
                'policy': policy_name,
                'n_runs': len(successful_runs),
            }

            # Aggregate test metrics (including MAPE)
            for metric in ['mae', 'mse', 'mape', 'r2']:
                values = [r['test_metrics'].get(metric, np.nan)
                         for r in successful_runs
                         if 'test_metrics' in r and metric in r['test_metrics']]
                if values:
                    metrics[f'{metric}_mean'] = np.mean(values)
                    metrics[f'{metric}_std'] = np.std(values)

            # Best expression (from best run based on R2)
            best_run = max(successful_runs,
                          key=lambda r: r.get('test_metrics', {}).get('r2', -np.inf))
            metrics['best_expression'] = best_run.get('best_expression', 'N/A')
            metrics['best_reward'] = best_run.get('best_reward', np.nan)

            # Collect all expressions from runs
            metrics['all_expressions'] = [r.get('best_expression', 'N/A') for r in successful_runs]

            summary_data.append(metrics)

        df = pd.DataFrame(summary_data)

        # Save summary
        summary_path = os.path.join(self.results_dir, 'summary.csv')
        df.to_csv(summary_path, index=False)
        print(f"\nSummary saved to {summary_path}")

        return df

    def print_detailed_results(self) -> None:
        """
        Print detailed results with symbolic expressions and metrics with std deviations.
        """
        if not self.results:
            print("No results to print. Run experiments first.")
            return

        print("\n" + "=" * 80)
        print("DETAILED EXPERIMENT RESULTS (10 runs per policy)")
        print("=" * 80)

        for policy_name, runs in self.results.items():
            successful_runs = [r for r in runs if r.get('success', False)]

            if not successful_runs:
                print(f"\n{policy_name.upper()}: No successful runs")
                continue

            print(f"\n{'=' * 80}")
            print(f"POLICY: {policy_name.upper()}")
            print(f"{'=' * 80}")
            print(f"Successful runs: {len(successful_runs)}/{len(runs)}")

            # Calculate metrics with std deviations
            metrics_summary = {}
            for metric in ['mae', 'mse', 'mape', 'r2']:
                values = [r['test_metrics'].get(metric, np.nan)
                         for r in successful_runs
                         if 'test_metrics' in r and metric in r['test_metrics']]
                if values:
                    metrics_summary[metric] = {
                        'mean': np.mean(values),
                        'std': np.std(values)
                    }

            # Print metrics
            print("\nTest Metrics (mean +/- std over 10 runs):")
            print("-" * 50)
            for metric, stats in metrics_summary.items():
                metric_name = metric.upper()
                if metric == 'mape':
                    print(f"  {metric_name}: {stats['mean']:.4f} +/- {stats['std']:.4f} %")
                else:
                    print(f"  {metric_name}: {stats['mean']:.4f} +/- {stats['std']:.4f}")

            # Find best expression (by R2)
            best_run = max(successful_runs,
                          key=lambda r: r.get('test_metrics', {}).get('r2', -np.inf))

            print("\nBest Symbolic Expression (highest R2):")
            print("-" * 50)
            print(f"  {best_run.get('best_expression', 'N/A')}")
            print(f"  (R2 = {best_run.get('test_metrics', {}).get('r2', 'N/A'):.4f})")

            # List all expressions found
            print("\nAll Expressions Found:")
            print("-" * 50)
            for i, run in enumerate(successful_runs):
                r2_val = run.get('test_metrics', {}).get('r2', np.nan)
                expr = run.get('best_expression', 'N/A')
                print(f"  Run {i+1} (R2={r2_val:.4f}): {expr}")

        # Save detailed results to file
        results_path = os.path.join(self.results_dir, 'detailed_results.txt')
        with open(results_path, 'w') as f:
            f.write("DETAILED EXPERIMENT RESULTS\n")
            f.write("=" * 80 + "\n\n")

            for policy_name, runs in self.results.items():
                successful_runs = [r for r in runs if r.get('success', False)]
                if not successful_runs:
                    continue

                f.write(f"POLICY: {policy_name.upper()}\n")
                f.write("-" * 40 + "\n")

                # Metrics
                for metric in ['mae', 'mse', 'mape', 'r2']:
                    values = [r['test_metrics'].get(metric, np.nan)
                             for r in successful_runs
                             if 'test_metrics' in r and metric in r['test_metrics']]
                    if values:
                        unit = ' %' if metric == 'mape' else ''
                        f.write(f"{metric.upper()}: {np.mean(values):.4f} +/- {np.std(values):.4f}{unit}\n")

                # Best expression
                best_run = max(successful_runs,
                              key=lambda r: r.get('test_metrics', {}).get('r2', -np.inf))
                f.write(f"\nBest Expression: {best_run.get('best_expression', 'N/A')}\n")
                f.write("\n" + "=" * 80 + "\n\n")

        print(f"\nDetailed results saved to {results_path}")

    def plot_results(self, save_path: Optional[str] = None) -> None:
        """
        Create visualization of experiment results.

        Args:
            save_path: Path to save the figure. If None, displays interactively.
        """
        if not self.results:
            print("No results to plot. Run experiments first.")
            return

        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        fig.suptitle('DSO Experiment Results: Policy Comparison (10 runs)', fontsize=14)

        metrics = ['mae', 'mse', 'mape', 'r2']
        metric_labels = ['MAE', 'MSE', 'MAPE (%)', 'R-squared']

        for ax, metric, label in zip(axes.flat, metrics, metric_labels):
            policy_names = []
            means = []
            stds = []

            for policy_name, runs in self.results.items():
                successful_runs = [r for r in runs if r.get('success', False)]
                values = [r['test_metrics'].get(metric, np.nan)
                         for r in successful_runs
                         if 'test_metrics' in r and metric in r['test_metrics']]

                if values:
                    policy_names.append(policy_name.replace('_', '\n'))
                    means.append(np.mean(values))
                    stds.append(np.std(values))

            x = np.arange(len(policy_names))
            bars = ax.bar(x, means, yerr=stds, capsize=5, alpha=0.7,
                         color=['#2ecc71', '#3498db', '#e74c3c'])
            ax.set_xticks(x)
            ax.set_xticklabels(policy_names)
            ax.set_ylabel(label)
            ax.set_title(f'{label} by Policy (mean +/- std)')
            ax.grid(axis='y', alpha=0.3)

            # Add value labels with std
            for bar, mean, std in zip(bars, means, stds):
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + std + 0.01,
                       f'{mean:.3f}\n+/-{std:.3f}', ha='center', va='bottom', fontsize=8)

        plt.tight_layout()

        if save_path is None:
            save_path = os.path.join(self.results_dir, 'policy_comparison.png')

        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Figure saved to {save_path}")
        plt.close()


# ============================================================================
# Command-line Interface
# ============================================================================

def run_dso_from_command_line(config_path: str, n_runs: int = 1, seed: int = 42) -> None:
    """
    Run DSO using command line (alternative to Python API).

    Args:
        config_path: Path to configuration JSON file
        n_runs: Number of independent runs
        seed: Starting random seed
    """
    import subprocess

    cmd = [
        'python', '-m', 'dso.run',
        config_path,
        f'--runs={n_runs}',
        f'--seed={seed}'
    ]

    print(f"Running: {' '.join(cmd)}")
    subprocess.run(cmd, check=True)


def main():
    """Main entry point for DSO experiments."""
    parser = argparse.ArgumentParser(
        description='Run Deep Symbolic Regression experiments with different RL policies'
    )
    parser.add_argument(
        '--policy',
        type=str,
        default='all',
        choices=['all', 'risk_seeking', 'vanilla', 'priority_queue'],
        help='Which policy to run (default: all)'
    )
    parser.add_argument(
        '--data-path',
        type=str,
        default='Propagation_models.xlsx',
        help='Path to Excel file with ABG data'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='dso_experiments',
        help='Directory for experiment outputs'
    )
    parser.add_argument(
        '--n-samples',
        type=int,
        default=50000,
        help='Number of samples per training run (default: 50000)'
    )
    parser.add_argument(
        '--batch-size',
        type=int,
        default=200,
        help='Batch size for training (default: 200)'
    )
    parser.add_argument(
        '--n-runs',
        type=int,
        default=10,
        help='Number of independent runs per policy (default: 10)'
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed'
    )
    parser.add_argument(
        '--dataset',
        type=str,
        default='abg',
        choices=['abg', 'ci', 'indoor', 'outdoor'],
        help='Dataset type: abg, ci, indoor, or outdoor (default: abg)'
    )
    parser.add_argument(
        '--prepare-only',
        action='store_true',
        help='Only prepare data and configs, do not run experiments'
    )

    args = parser.parse_args()

    # Set default data path based on dataset type
    if args.data_path == 'Propagation_models.xlsx':
        if args.dataset in ['indoor', 'outdoor']:
            args.data_path = 'csv_experiments.xlsx'

    # Determine which policies to run
    if args.policy == 'all':
        policies = ['risk_seeking', 'vanilla', 'priority_queue']
    else:
        policies = [args.policy]

    # Set output directory to include dataset type if using default
    output_dir = args.output_dir
    if output_dir == 'dso_experiments':
        output_dir = f'dso_experiments_{args.dataset}'

    # Get policy parameters for the selected dataset
    dataset_policy_params = DSOExperiment.DATASET_POLICY_PARAMS.get(args.dataset, DSOExperiment.ABG_POLICY_PARAMS)

    print("=" * 80)
    print(f"Deep Symbolic Regression Experiments ({args.dataset.upper()} dataset)")
    print("=" * 80)
    print(f"Dataset: {args.dataset.upper()}")
    print(f"Policies: {', '.join(policies)}")
    print(f"Data path: {args.data_path}")
    print(f"Output dir: {output_dir}")
    print(f"Samples per run: {args.n_samples}")
    print(f"Number of runs: {args.n_runs}")
    print(f"Random seed: {args.seed}")
    print(f"\nPolicy-specific parameters ({args.dataset.upper()}):")
    for policy_name, params in dataset_policy_params.items():
        print(f"  {policy_name}: lr={params['learning_rate']}, entropy={params['entropy_weight']}, batch={params['batch_size']}")
    print("\nConstraints:")
    print("  - Token length: 4-40")
    print("  - No all-constant children")
    print("  - No inverse of unary ops as children")
    print("  - No nested trig operators")
    print("  - Repeat constraint on f and d")
    print("=" * 80)

    # Create experiment manager
    experiment = DSOExperiment(
        data_path=args.data_path,
        output_dir=output_dir,
        n_samples=args.n_samples,
        batch_size=args.batch_size,
        n_runs=args.n_runs,
        random_seed=args.seed,
        dataset_type=args.dataset
    )

    # Prepare data
    experiment.prepare_data()

    # Create configs
    experiment.create_configs(policies)

    if args.prepare_only:
        print("\nData and configs prepared. Use --prepare-only=False to run experiments.")
        return

    # Run experiments
    experiment.run_experiments(policies)

    # Print detailed results with expressions and metrics
    experiment.print_detailed_results()

    # Create summary table
    print("\n" + "=" * 80)
    print("SUMMARY TABLE")
    print("=" * 80)

    summary = experiment.summarize_results()
    if not summary.empty:
        # Print a clean summary table
        for _, row in summary.iterrows():
            print(f"\n{row['policy'].upper()}:")
            print(f"  MAE:  {row.get('mae_mean', np.nan):.4f} +/- {row.get('mae_std', np.nan):.4f}")
            print(f"  MSE:  {row.get('mse_mean', np.nan):.4f} +/- {row.get('mse_std', np.nan):.4f}")
            print(f"  MAPE: {row.get('mape_mean', np.nan):.4f} +/- {row.get('mape_std', np.nan):.4f} %")
            print(f"  R2:   {row.get('r2_mean', np.nan):.4f} +/- {row.get('r2_std', np.nan):.4f}")
            print(f"  Best Expression: {row.get('best_expression', 'N/A')}")

    # Create visualization
    experiment.plot_results()

    print("\n" + "=" * 80)
    print("Experiments complete!")
    print(f"Results saved to: {output_dir}")
    print("  - run_metrics.csv: Per-run metrics (MAE, MSE, MAPE, R2) for averaging")
    print("  - detailed_results.txt: Symbolic expressions and metrics")
    print("  - summary.csv: Summary statistics")
    print("  - policy_comparison.png: Visualization")
    print("  - experiment_results.json: Raw results")
    print("=" * 80)


def print_parallel_commands():
    """Print commands for running experiments in parallel terminals."""
    print("\n" + "=" * 80)
    print("COMMANDS FOR PARALLEL TERMINAL EXECUTION")
    print("=" * 80)

    print("\n--- ABG Dataset ---")
    print("Open 3 terminal windows and run one command in each:\n")
    print("# Terminal 1 - ABG Risk-Seeking (lr=0.002, entropy=0.008):")
    print("python dso_experiments.py --dataset abg --policy risk_seeking --output-dir dso_experiments_abg_50k_risk_seeking\n")
    print("# Terminal 2 - ABG Vanilla PG (lr=0.0001, entropy=0.005):")
    print("python dso_experiments.py --dataset abg --policy vanilla --output-dir dso_experiments_abg_50k_vanilla\n")
    print("# Terminal 3 - ABG Priority Queue (lr=0.002, entropy=0.005):")
    print("python dso_experiments.py --dataset abg --policy priority_queue --output-dir dso_experiments_abg_50k_priority_queue\n")

    print("\n--- CI Dataset ---")
    print("Open 3 terminal windows and run one command in each:\n")
    print("# Terminal 1 - CI Risk-Seeking (lr=0.001, entropy=0.008):")
    print("python dso_experiments.py --dataset ci --policy risk_seeking --output-dir dso_experiments_ci_50k_risk_seeking\n")
    print("# Terminal 2 - CI Vanilla PG (lr=0.0005, entropy=0.008):")
    print("python dso_experiments.py --dataset ci --policy vanilla --output-dir dso_experiments_ci_50k_vanilla\n")
    print("# Terminal 3 - CI Priority Queue (lr=0.002, entropy=0.005):")
    print("python dso_experiments.py --dataset ci --policy priority_queue --output-dir dso_experiments_ci_50k_priority_queue\n")

    print("\n--- Indoor Dataset ---")
    print("Open 3 terminal windows and run one command in each:\n")
    print("# Terminal 1 - Indoor Risk-Seeking (lr=0.0005, entropy=0.03):")
    print("python dso_experiments.py --dataset indoor --policy risk_seeking --output-dir dso_experiments_indoor_50k_risk_seeking\n")
    print("# Terminal 2 - Indoor Vanilla PG (lr=0.001, entropy=0.02):")
    print("python dso_experiments.py --dataset indoor --policy vanilla --output-dir dso_experiments_indoor_50k_vanilla\n")
    print("# Terminal 3 - Indoor Priority Queue (lr=0.001, entropy=0.01):")
    print("python dso_experiments.py --dataset indoor --policy priority_queue --output-dir dso_experiments_indoor_50k_priority_queue\n")

    print("\n--- Outdoor Dataset ---")
    print("Open 3 terminal windows and run one command in each:\n")
    print("# Terminal 1 - Outdoor Risk-Seeking (lr=0.0005, entropy=0.01):")
    print("python dso_experiments.py --dataset outdoor --policy risk_seeking --output-dir dso_experiments_outdoor_50k_risk_seeking\n")
    print("# Terminal 2 - Outdoor Vanilla PG (lr=0.0001, entropy=0.01):")
    print("python dso_experiments.py --dataset outdoor --policy vanilla --output-dir dso_experiments_outdoor_50k_vanilla\n")
    print("# Terminal 3 - Outdoor Priority Queue (lr=0.0005, entropy=0.01):")
    print("python dso_experiments.py --dataset outdoor --policy priority_queue --output-dir dso_experiments_outdoor_50k_priority_queue\n")
    print("=" * 80)


if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == '--show-parallel-commands':
        print_parallel_commands()
    else:
        main()