2025-01-03 20:10:22,109 - INFO - TrainLoopRank1 - Rank 1: Collected 10000 steps.   is it possible to process all of the steps sequence before  because producing steps sequence at every epochs is taking a lot of time and increase the training time.  we need to preprocess multiple sequence step at first, save them and load them for training the model at every epochs  , here are the scripts for reference : {
    "./main.py": "# File: main.py\n\nimport os\nimport logging\nimport torch\nimport torch.multiprocessing as mp\nimport yaml\n\nfrom training.runner import train_worker\nfrom training.utils import (\n    validate_and_cast_config,\n    setup_logging,\n    setup_seed,\n)\n\n\ndef main(config):\n    \"\"\"\n    Entry point for training. Initializes multiprocessing for distributed training.\n    \"\"\"\n    logger = logging.getLogger(\"Main\")\n    logger.info(\"Entering main() function.\")\n\n    # Set random seed\n    setup_seed(config[\"random_seed\"])\n\n    # Number of GPUs\n    num_gpus = torch.cuda.device_count()\n    world_size = num_gpus if num_gpus >= 1 else 1\n    logger.info(f\"Number of GPUs available: {num_gpus}\")\n\n    if num_gpus > 1:\n        # Use DistributedDataParallel (DDP)\n        logger.info(\"Using DistributedDataParallel for multi-GPU training.\")\n\n        # Set up environment variables for distributed training\n        os.environ[\"MASTER_ADDR\"] = config[\"distributed\"][\"master_addr\"]\n        os.environ[\"MASTER_PORT\"] = config[\"distributed\"][\"master_port\"]\n        logger.info(f\"MASTER_ADDR set to {config['distributed']['master_addr']}\")\n        logger.info(f\"MASTER_PORT set to {config['distributed']['master_port']}\")\n\n        try:\n            mp.spawn(\n                train_worker,\n                args=(config, world_size),\n                nprocs=world_size,\n                join=True\n            )\n        except Exception as e:\n            logger.error(f\"Training failed: {e}\", exc_info=True)\n            raise\n    else:\n        # Single GPU or CPU\n        logger.info(\"Single GPU or CPU detected.\")\n        train_worker(0, config, world_size)\n\n\nif __name__ == \"__main__\":\n    logger = setup_logging(\"main.log\")\n    logger.info(\"Starting main execution.\")\n\n    config_path = \"config/config.yaml\"\n    if not os.path.exists(config_path):\n        logger.error(f\"Configuration file not found at path: {config_path}\")\n        raise FileNotFoundError(f\"Configuration file not found at path: {config_path}\")\n\n    try:\n        with open(config_path, \"r\") as f:\n            raw_cfg = yaml.safe_load(f)\n            logger.debug(f\"Loaded raw config: {raw_cfg}\")\n    except yaml.YAMLError as e:\n        logger.error(f\"Failed to parse the configuration file: {e}\")\n        raise\n\n    try:\n        config = validate_and_cast_config(raw_cfg)\n        logger.info(\"Configuration validated successfully.\")\n        logger.info(f\"Loaded configuration:\\n{yaml.dump(config)}\")  # Log the full configuration\n    except Exception as e:\n        logger.error(f\"Configuration validation failed: {e}\", exc_info=True)\n        raise\n\n    logger.info(\"Configuration loaded and validated successfully.\")\n\n    # Start\n    try:\n        main(config)\n    except Exception as e:\n        logger.error(f\"Training failed with exception: {e}\", exc_info=True)\n        raise",
    "./collection_json.py": "# File: collection_json.py\n\nimport os\nimport json\nimport argparse\nimport logging\n\ndef setup_logging():\n    \"\"\"\n    Configure logging to output to both file and console for this script.\n    \"\"\"\n    logger = logging.getLogger(\"CollectionJSON\")\n    logger.setLevel(logging.DEBUG)\n\n    # Create console handler\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n\n    # Create file handler\n    fh = logging.FileHandler(\"logs/collection_json.log\")\n    fh.setLevel(logging.DEBUG)\n\n    # Create formatter and add to handlers\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(message)s\")\n    ch.setFormatter(formatter)\n    fh.setFormatter(formatter)\n\n    # Add handlers to the logger if none exist\n    if not logger.hasHandlers():\n        logger.addHandler(ch)\n        logger.addHandler(fh)\n\n    return logger\n\ndef collect_scripts_to_json(root_dir, output_file, logger):\n    \"\"\"\n    Traverse the given directory and its subdirectories to collect .py and .yaml files,\n    then save their contents into a JSON file.\n    \"\"\"\n    logger.info(\"Starting script collection.\")\n    scripts = {}\n\n    for dirpath, _, filenames in os.walk(root_dir):\n        for file in filenames:\n            if file.endswith(('.py', '.yaml')):\n                file_path = os.path.join(dirpath, file)\n                try:\n                    with open(file_path, 'r') as f:\n                        scripts[file_path] = f.read()\n                    logger.debug(f\"Collected script: {file_path}\")\n                except Exception as e:\n                    logger.error(f\"Failed to read {file_path}: {e}\")\n\n    try:\n        with open(output_file, 'w') as json_file:\n            json.dump(scripts, json_file, indent=4)\n        logger.info(f\"Scripts successfully collected and saved to {output_file}\")\n    except Exception as e:\n        logger.error(f\"Failed to write to {output_file}: {e}\")\n\ndef main():\n    parser = argparse.ArgumentParser(description=\"Collect .py and .yaml scripts into a JSON file.\")\n    parser.add_argument('--root_dir', type=str, default='./', help='Root directory to search for scripts.')\n    parser.add_argument('--output_file', type=str, default='scripts_collection.json', help='Output JSON file path.')\n\n    args = parser.parse_args()\n\n    logger = setup_logging()\n    logger.info(f\"Starting collection from root directory: {args.root_dir}\")\n\n    collect_scripts_to_json(args.root_dir, args.output_file, logger)\n\nif __name__ == \"__main__\":\n    main()",
    "./monitoring.py": "# File: monitoring.py\n\nimport os\nimport logging\nimport json\nfrom torch.utils.tensorboard import SummaryWriter\nfrom training.utils import log_memory_usage, log_gpu_memory\n\nlogger = logging.getLogger(\"Monitoring\")\nlogger.setLevel(logging.DEBUG)\nif not logger.hasHandlers():\n    os.makedirs(\"logs\", exist_ok=True)\n    fh = logging.FileHandler(\"logs/monitoring.log\")\n    fh.setLevel(logging.DEBUG)\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n\ndef monitor_training(\n    epoch: int,\n    rank: int,\n    avg_metrics: dict,\n    eval_metrics: dict,\n    writer: SummaryWriter = None\n) -> None:\n    \"\"\"\n    Monitor and log training progress, metrics, and insights.\n\n    Args:\n        epoch (int): Current epoch number.\n        rank (int): Process rank (for compatibility).\n        avg_metrics (dict): Dictionary of average training metrics.\n        eval_metrics (dict): Dictionary of evaluation metrics.\n        writer (SummaryWriter, optional): TensorBoard SummaryWriter.\n    \"\"\"\n    logs_dir = \"logs\"\n    os.makedirs(logs_dir, exist_ok=True)\n\n    # Log metrics in console logs\n    logger.info(f\"Epoch {epoch + 1} - Training Metrics: {avg_metrics}\")\n    logger.info(f\"Epoch {epoch + 1} - Evaluation Metrics: {eval_metrics}\")\n\n    # Write metrics to a JSON file\n    metrics_file = os.path.join(logs_dir, \"metrics.json\")\n    try:\n        with open(metrics_file, \"a\") as f:\n            json.dump(\n                {\n                    \"epoch\": epoch + 1,\n                    \"train_metrics\": avg_metrics,\n                    \"eval_metrics\": eval_metrics,\n                },\n                f,\n            )\n            f.write(\"\\n\")\n        logger.info(f\"Metrics appended to {metrics_file}\")\n    except Exception as e:\n        logger.error(f\"Failed to write to {metrics_file}: {e}\")\n\n    if writer:\n        # Write training metrics\n        for k, v in avg_metrics.items():\n            writer.add_scalar(f\"Train/{k}\", v, epoch)\n        # Write evaluation metrics\n        for k, v in eval_metrics.items():\n            writer.add_scalar(f\"Eval/{k}\", v, epoch)\n        writer.flush()\n\n    # Optionally log CPU/GPU usage\n    log_memory_usage(logger, phase=\"After Logging Metrics\")\n    log_gpu_memory(logger, phase=\"After Logging Metrics\")\n\n    # Interpret metrics\n    interpret_metrics(avg_metrics, eval_metrics, logger)\n\n\ndef interpret_metrics(avg_metrics: dict, eval_metrics: dict, logger: logging.Logger) -> None:\n    \"\"\"\n    Interpret and extrapolate insights from metrics.\n\n    Args:\n        avg_metrics (dict): Average training metrics.\n        eval_metrics (dict): Evaluation metrics.\n        logger (logging.Logger): Logger instance.\n    \"\"\"\n    # Example interpretations\n    if \"actor_loss\" in avg_metrics and \"critic_loss\" in avg_metrics:\n        logger.info(f\"Actor Loss: {avg_metrics['actor_loss']:.6f}\")\n        logger.info(f\"Critic Loss: {avg_metrics['critic_loss']:.6f}\")\n\n    if \"entropy\" in avg_metrics:\n        if avg_metrics[\"entropy\"] > 0.01:\n            logger.info(\"Model is exploring adequately.\")\n        else:\n            logger.info(\"Entropy is low; model may be exploiting too early.\")\n\n    # Evaluation insights\n    if \"average_reward\" in eval_metrics:\n        logger.info(f\"Average Reward: {eval_metrics['average_reward']:.6f}\")\n    if \"sharpe_ratio\" in eval_metrics:\n        logger.info(f\"Sharpe Ratio: {eval_metrics['sharpe_ratio']:.6f}\")\n    if \"max_drawdown\" in eval_metrics:\n        logger.info(f\"Max Drawdown: {eval_metrics['max_drawdown']:.6f}\")\n    if \"average_win_rate\" in eval_metrics:\n        logger.info(f\"Average Win Rate: {eval_metrics['average_win_rate']:.2%}\")\n\n\ndef setup_tensorboard(rank: int, log_dir: str = \"runs\") -> SummaryWriter:\n    \"\"\"\n    Setup TensorBoard writer.\n\n    Args:\n        rank (int): Process rank.\n        log_dir (str): Directory for TensorBoard logs.\n\n    Returns:\n        SummaryWriter: TensorBoard SummaryWriter instance.\n    \"\"\"\n    writer = None\n    os.makedirs(log_dir, exist_ok=True)\n    writer = SummaryWriter(log_dir=log_dir)\n    logger.info(f\"TensorBoard writer set up at {log_dir}.\")\n\n    return writer",
    "./data/test.py": "import pandas as pd\nimport torch\nimport os\n\ndef scan_csv_for_nan(file_path):\n    \"\"\"\n    Scans a CSV file for NaN values or empty cases.\n\n    :param file_path: Path to the CSV file.\n    \"\"\"\n    if not os.path.exists(file_path):\n        print(f\"File not found: {file_path}\")\n        return\n\n    try:\n        df = pd.read_csv(file_path)\n        print(f\"Scanning CSV file: {file_path}\")\n        if df.empty:\n            print(\"The CSV file is empty.\")\n        else:\n            nan_count = df.isna().sum().sum()\n            if nan_count > 0:\n                print(f\"The CSV file contains {nan_count} NaN values.\")\n            else:\n                print(\"No NaN values found in the CSV file.\")\n    except Exception as e:\n        print(f\"Error reading CSV file: {e}\")\n\ndef scan_pt_for_nan(file_path):\n    \"\"\"\n    Scans a PyTorch .pt file for NaN values or empty tensors.\n\n    :param file_path: Path to the .pt file.\n    \"\"\"\n    if not os.path.exists(file_path):\n        print(f\"File not found: {file_path}\")\n        return\n\n    try:\n        data = torch.load(file_path)\n        print(f\"Scanning .pt file: {file_path}\")\n\n        if isinstance(data, torch.Tensor):\n            if data.numel() == 0:\n                print(\"The .pt file contains an empty tensor.\")\n            elif torch.isnan(data).any():\n                print(\"The .pt file contains NaN values in the tensor.\")\n            else:\n                print(\"No NaN values found in the .pt file.\")\n        elif isinstance(data, dict):\n            nan_count = 0\n            empty_tensors = 0\n            for key, tensor in data.items():\n                if isinstance(tensor, torch.Tensor):\n                    if tensor.numel() == 0:\n                        empty_tensors += 1\n                        print(f\"Empty tensor found for key: {key}\")\n                    elif torch.isnan(tensor).any():\n                        nan_count += torch.isnan(tensor).sum().item()\n                        print(f\"NaN values found for key: {key}\")\n            if nan_count == 0 and empty_tensors == 0:\n                print(\"No NaN or empty tensors found in the .pt file.\")\n        else:\n            print(\"The .pt file contains unrecognized data.\")\n    except Exception as e:\n        print(f\"Error reading .pt file: {e}\")\n\n# Paths to your files\ncsv_file_path = \"dataset.csv\"  # Replace with your CSV file path\npt_file_path = \"preprocessed_data.pt\"  # Replace with your .pt file path\n\n# Perform scans\nscan_csv_for_nan(csv_file_path)\nscan_pt_for_nan(pt_file_path)",
    "./environment/trading_env.py": "# File: environment/trading_env.py\n\nimport gym\nfrom gym import spaces\nimport torch\nimport numpy as np\nimport os\nimport glob\nimport random\nimport pandas as pd\n\n\nclass TradingEnv(gym.Env):\n    \"\"\"\n    Custom Trading Environment that follows OpenAI Gym interface.\n    \"\"\"\n    metadata = {'render.modes': ['human']}\n\n    def __init__(self, config):\n        super(TradingEnv, self).__init__()\n\n        self.config = config\n\n        # Initialize attributes\n        self.feature_columns = config['data']['feature_columns']\n        self.n_steps = config['data']['n_steps']\n        self.max_steps = config['environment']['max_steps']\n\n        # Load data from PT files\n        pt_data_dir = config['data'].get('pt_data_dir', 'data/processed/')\n        pt_data_prefix = config['data'].get('pt_data_prefix', 'processed_data_part')\n        pt_files = sorted(glob.glob(os.path.join(pt_data_dir, f\"{pt_data_prefix}_*.pt\")))\n\n        if not pt_files:\n            raise FileNotFoundError(f\"No processed data files found in {pt_data_dir}\")\n\n        self.pt_files = pt_files\n        self.current_pt_file_idx = -1\n        self.data = None\n\n        # Load raw close prices\n        raw_data_path = config['data']['raw_data_path']\n        self.raw_df = pd.read_csv(raw_data_path)\n        self.close_prices = self.raw_df['close'].values\n\n        # Call load_next_data_chunk() after initializing attributes\n        self.load_next_data_chunk()\n\n        # Define action and observation space\n        self.action_space = spaces.Discrete(config['model']['action_dim'])\n        self.observation_space = spaces.Box(\n            low=-np.inf,\n            high=np.inf,\n            shape=(self.n_steps, len(self.feature_columns)),\n            dtype=np.float32\n        )\n\n        # Initialize variables\n        self.current_step = 0\n        self.position = 0\n        self.total_profit = 0.0\n\n    def load_next_data_chunk(self):\n        self.current_pt_file_idx = (self.current_pt_file_idx + 1) % len(self.pt_files)\n        self.data = torch.load(self.pt_files[self.current_pt_file_idx])\n\n        # Recalculate end_step based on new data\n        self.end_step = len(self.data) - self.n_steps\n        self.current_step = 0  # Reset current_step when loading new data\n\n    def reset(self):\n        if self.current_step >= self.end_step:\n            self.load_next_data_chunk()\n\n        self.current_step = 0\n        self.position = 0\n        self.total_profit = 0.0\n        self.entry_price = None\n        return self._next_observation()\n\n    def _next_observation(self):\n        obs = self.data[self.current_step:self.current_step + self.n_steps]\n        if torch.isnan(obs).any():\n            raise ValueError(f\"Observation contains NaN values at step {self.current_step}\")\n        return obs.numpy().astype(np.float32)\n\n    def step(self, action):\n        idx = self.current_step + self.n_steps - 1\n        if idx >= len(self.data):\n            done = True\n            obs = np.zeros_like(self.observation_space.sample())\n            info = {'total_profit': self.total_profit}\n            return obs, 0.0, done, info\n\n        current_price = self.close_prices[idx]  # Use the unnormalized close price\n\n        reward = 0.0\n\n        # Implement trading logic\n        if action == 0:  # Sell\n            if self.position == 1:  # If long, close position\n                reward = current_price - self.entry_price\n                self.total_profit += reward\n                self.position = 0\n                self.entry_price = None\n            elif self.position == 0:  # Open short position\n                self.position = -1\n                self.entry_price = current_price\n        elif action == 2:  # Buy\n            if self.position == -1:  # If short, close position\n                reward = self.entry_price - current_price\n                self.total_profit += reward\n                self.position = 0\n                self.entry_price = None\n            elif self.position == 0:  # Open long position\n                self.position = 1\n                self.entry_price = current_price\n        # If action == 1 (Hold), do nothing\n\n        self.current_step += 1\n        done = self.current_step >= self.end_step or self.current_step >= self.max_steps\n\n        obs = self._next_observation() if not done else np.zeros_like(self.observation_space.sample())\n\n        info = {'total_profit': self.total_profit}\n\n        return obs, reward, done, info\n\n    def render(self, mode='human'):\n        pass",
    "./training/runner.py": "# File: training/runner.py\n\nimport os\nimport torch\nimport torch.distributed as dist\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom training.train_loop import train_model\nfrom training.utils import setup_logging, setup_seed, load_checkpoint\nfrom environment.trading_env import TradingEnv\nfrom models.models import PPOAgent\nfrom data_processing.preprocess import export_and_save_preprocessed_data\nimport logging\n\n\ndef train_worker(rank, config, world_size):\n    \"\"\"\n    Train worker function for distributed training with DDP.\n    \"\"\"\n    logger = setup_logging(f\"runner_rank{rank}.log\")\n    logger.info(f\"Starting training worker on rank {rank}.\")\n\n    # Set random seed for each process\n    setup_seed(config.get(\"random_seed\", 42) + rank)\n    logger.info(f\"Random seed initialized for Rank {rank}.\")\n\n    # Initialize process group\n    dist.init_process_group(\n        backend=\"nccl\",\n        init_method=\"env://\",\n        world_size=world_size,\n        rank=rank\n    )\n\n    # Map rank to GPU id (ensure that GPU ids are correct)\n    torch.cuda.set_device(rank)\n    device = torch.device(f\"cuda:{rank}\")\n\n    logger.info(f\"Process group initialized for Rank {rank}. Device: {device}\")\n\n    # Only rank 0 performs data preprocessing\n    if rank == 0:\n        logger.info(\"Rank 0: Starting data preprocessing.\")\n        export_and_save_preprocessed_data(config, logger)\n    else:\n        logger.info(f\"Rank {rank}: Waiting for data preprocessing to complete.\")\n\n    # Ensure all processes wait for data preprocessing to complete\n    dist.barrier()\n\n    # Initialize environment\n    try:\n        env = TradingEnv(config)\n        logger.info(f\"Environment initialized for Rank {rank}.\")\n    except Exception as e:\n        logger.error(f\"Failed to initialize environment: {e}\", exc_info=True)\n        dist.destroy_process_group()\n        raise\n\n    # Initialize model\n    try:\n        model = PPOAgent(\n            feature_dim=config[\"model\"][\"feature_dim\"],\n            hidden_dim=config[\"model\"][\"ppo_hidden_dim\"],\n            action_dim=config[\"model\"][\"action_dim\"],\n            n_heads=config[\"model\"][\"ppo_n_heads\"],\n            n_layers=config[\"model\"][\"ppo_n_layers\"],\n            dropout=config[\"model\"][\"ppo_dropout\"],\n            temperature=config[\"model\"].get('temperature', 1.0),\n            layer_norm=config[\"model\"].get('layer_norm', False)\n        ).to(device)\n        logger.info(f\"PPOAgent initialized on Rank {rank}.\")\n    except Exception as e:\n        logger.error(f\"Failed to initialize PPOAgent: {e}\", exc_info=True)\n        dist.destroy_process_group()\n        raise\n\n    # Load checkpoint if required\n    start_epoch = 0\n    best_reward = float(\"-inf\")\n    resume_ckpt = config[\"checkpoint\"].get(\"resume_checkpoint\", \"\")\n    if resume_ckpt:\n        logger.info(f\"resume_checkpoint from config => '{resume_ckpt}'\")\n        if os.path.exists(resume_ckpt):\n            logger.info(f\"Found checkpoint '{resume_ckpt}'. Attempting to load...\")\n            epoch_loaded, best_reward_loaded = load_checkpoint(\n                resume_ckpt,\n                model,\n                None,  # If you want to also load optimizer states, pass them here\n                None,\n                device,\n            )\n            start_epoch = epoch_loaded + 1\n            best_reward = best_reward_loaded\n            logger.info(f\"Resuming from epoch={start_epoch}, best_reward={best_reward}\")\n        else:\n            logger.warning(f\"resume_checkpoint set but file not found => '{resume_ckpt}'. Starting fresh.\")\n    else:\n        logger.info(\"No 'resume_checkpoint' specified in config. Starting from epoch 0.\")\n\n    # Wrap model in DDP\n    model = DDP(model, device_ids=[rank], output_device=rank)\n\n    try:\n        logger.info(f\"Starting the training loop on Rank {rank}.\")\n        # Pass start_epoch and best_reward to train_model\n        train_model(\n            env=env,\n            config=config,\n            device=device,\n            model=model,\n            rank=rank,\n            world_size=world_size,\n            start_epoch=start_epoch,\n            num_epochs=config[\"training\"][\"epochs\"],\n            best_reward=best_reward\n        )\n    except Exception as e:\n        logger.error(f\"Training failed with exception on Rank {rank}: {e}\", exc_info=True)\n    finally:\n        dist.destroy_process_group()\n        logger.info(f\"Training worker on Rank {rank} has completed.\")",
    "./training/train_loop.py": "# File: training/train_loop.py\n\nimport os\nimport logging\nimport torch\nimport numpy as np\nimport torch.distributed as dist  # Distributed training\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom training.utils import (\n    compute_advantages_and_returns,\n    save_checkpoint,\n    log_memory_usage,\n    log_gpu_memory,\n)\nfrom monitoring import monitor_training, setup_tensorboard\nfrom training.evaluation import evaluate_model\nfrom torch.distributions import Categorical\nimport gc\n\n# Import autocast and GradScaler from torch.cuda.amp\nfrom torch.cuda.amp import autocast, GradScaler\n\ndef collect_trajectories(env, model, device, config):\n    \"\"\"\n    Collect trajectories by interacting with the environment.\n    \"\"\"\n    model.eval()\n    trajectories = {\n        'states': [],\n        'actions': [],\n        'rewards': [],\n        'dones': [],\n        'values': [],\n        'log_probs': []\n    }\n    state = env.reset()\n    done = False\n    steps = 0\n    max_steps_per_episode = config[\"environment\"][\"max_steps\"]\n\n    while not done and steps < max_steps_per_episode:\n        state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(device)\n        with torch.no_grad():\n            if device.type == 'cuda':\n                with autocast():\n                    logits, value = model(state_tensor)\n            else:\n                logits, value = model(state_tensor)\n\n            action_probs = torch.softmax(logits, dim=-1)\n            action_dist = Categorical(probs=action_probs)\n            action = action_dist.sample()\n            log_prob = action_dist.log_prob(action)\n            value = value.item()\n\n        next_state, reward, done, _ = env.step(action.item())\n\n        # Store experience\n        trajectories['states'].append(state)\n        trajectories['actions'].append(action.item())\n        trajectories['rewards'].append(reward)\n        trajectories['dones'].append(done)\n        trajectories['values'].append(value)\n        trajectories['log_probs'].append(log_prob.item())\n\n        state = next_state\n        steps += 1\n\n    model.train()\n    return trajectories\n\ndef train_model(\n    env,\n    config,\n    device,\n    model,\n    rank,\n    world_size,\n    start_epoch,\n    num_epochs,\n    best_reward=-float(\"inf\"),\n):\n    \"\"\"\n    Main training loop for PPO with DDP.\n    \"\"\"\n    logger = logging.getLogger(f\"TrainLoopRank{rank}\")\n    logger.setLevel(logging.DEBUG)\n\n    # Setup logger handlers if not already set\n    if not logger.hasHandlers():\n        # Setup logger\n        os.makedirs(config['logging']['log_dir'], exist_ok=True)\n        fh = logging.FileHandler(f\"{config['logging']['log_dir']}/train_loop_rank{rank}.log\")\n        fh.setLevel(logging.DEBUG)\n        ch = logging.StreamHandler()\n        ch.setLevel(logging.INFO)\n        fmt = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n        fh.setFormatter(fmt)\n        ch.setFormatter(fmt)\n        logger.addHandler(fh)\n        logger.addHandler(ch)\n\n    logger.info(f\"Starting training from epoch {start_epoch + 1} to {num_epochs}\")\n    logger.info(f\"Rank {rank}: Using batch size {config['training']['batch_size']}\")\n\n    # Optimizer setup\n    optimizer = torch.optim.AdamW(\n        model.parameters(),\n        lr=config[\"training\"][\"actor_lr\"],  # Assuming same LR for both actor and critic\n        weight_decay=config[\"training\"].get(\"weight_decay\", 0.0),\n    )\n\n    max_grad_norm = config[\"training\"].get(\"max_grad_norm\", 0.5)\n    batch_size = config[\"training\"][\"batch_size\"]\n    ppo_epochs = config[\"training\"][\"ppo_epochs\"]\n\n    # Setup TensorBoard writer (only on rank 0)\n    writer = None\n    if config[\"logging\"][\"use_tensorboard\"] and rank == 0:\n        writer = setup_tensorboard(rank, log_dir=config['logging']['log_dir'])\n\n    # Initialize AMP scaler\n    scaler = GradScaler()\n\n    for epoch in range(start_epoch, num_epochs):\n        logger.info(f\"Rank {rank}, Epoch {epoch + 1}/{num_epochs}\")\n        log_memory_usage(logger, phase=f\"Epoch {epoch + 1} Start\")\n        log_gpu_memory(logger, phase=f\"Epoch {epoch + 1} Start\")\n\n        # Collect trajectories\n        trajectories = collect_trajectories(env, model, device, config)\n        logger.info(f\"Rank {rank}: Collected {len(trajectories['rewards'])} steps.\")\n\n        # Convert collected data to tensors\n        states = torch.stack([torch.from_numpy(np.array(s)) for s in trajectories['states']]).float().to(device)\n        actions = torch.tensor(trajectories['actions'], dtype=torch.long, device=device)\n        rewards = torch.tensor(trajectories['rewards'], dtype=torch.float32, device=device)\n        dones = torch.tensor(trajectories['dones'], dtype=torch.float32, device=device)\n        values = torch.tensor(trajectories['values'], dtype=torch.float32, device=device)\n        log_probs_old = torch.tensor(trajectories['log_probs'], dtype=torch.float32, device=device)\n\n        # Compute advantages and returns\n        advantages, returns = compute_advantages_and_returns(\n            rewards=rewards,\n            values=values,\n            dones=dones,\n            gamma=config[\"training\"][\"gamma\"],\n            gae_lambda=config[\"training\"][\"gae_lambda\"],\n            reward_scale=config['data'].get('reward_scale', 1.0)\n        )\n\n        # Normalize advantages\n        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)\n\n        # PPO Update\n        for ppo_epoch in range(ppo_epochs):\n            # Shuffle the data\n            indices = torch.randperm(len(states))\n            # Now split indices into batches\n            for start_idx in range(0, len(states), batch_size):\n                end_idx = start_idx + batch_size\n                batch_indices = indices[start_idx:end_idx]\n\n                batch_states = states[batch_indices]\n                batch_actions = actions[batch_indices]\n                batch_returns = returns[batch_indices]\n                batch_advantages = advantages[batch_indices]\n                batch_log_probs_old = log_probs_old[batch_indices]\n\n                logger.debug(f\"Rank {rank}, Epoch {epoch + 1}, Batch size: {batch_states.size(0)}\")\n\n                optimizer.zero_grad()\n                if device.type == 'cuda':\n                    with autocast():\n                        logits, values_pred = model(batch_states)\n                        action_dist = Categorical(logits=logits)\n                        log_probs = action_dist.log_prob(batch_actions)\n                        entropy = action_dist.entropy().mean()\n\n                        # Calculate ratios\n                        ratios = torch.exp(log_probs - batch_log_probs_old)\n\n                        # Surrogate loss\n                        surr1 = ratios * batch_advantages\n                        surr2 = torch.clamp(ratios, 1.0 - config[\"training\"][\"clip_epsilon\"],\n                                            1.0 + config[\"training\"][\"clip_epsilon\"]) * batch_advantages\n                        actor_loss = -torch.min(surr1, surr2).mean() - config[\"training\"][\"entropy_coeff\"] * entropy\n\n                        # Critic loss\n                        critic_loss = torch.nn.functional.mse_loss(values_pred.squeeze(-1), batch_returns)\n\n                        # Total loss\n                        total_loss = actor_loss + config[\"training\"][\"critic_coeff\"] * critic_loss\n\n                    # Backward pass with AMP\n                    scaler.scale(total_loss).backward()\n\n                    # Gradient clipping\n                    scaler.unscale_(optimizer)\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n\n                    scaler.step(optimizer)\n                    scaler.update()\n                else:\n                    # Forward pass without autocast\n                    logits, values_pred = model(batch_states)\n                    action_dist = Categorical(logits=logits)\n                    log_probs = action_dist.log_prob(batch_actions)\n                    entropy = action_dist.entropy().mean()\n\n                    # Calculate ratios\n                    ratios = torch.exp(log_probs - batch_log_probs_old)\n\n                    # Surrogate loss\n                    surr1 = ratios * batch_advantages\n                    surr2 = torch.clamp(ratios, 1.0 - config[\"training\"][\"clip_epsilon\"],\n                                        1.0 + config[\"training\"][\"clip_epsilon\"]) * batch_advantages\n                    actor_loss = -torch.min(surr1, surr2).mean() - config[\"training\"][\"entropy_coeff\"] * entropy\n\n                    # Critic loss\n                    critic_loss = torch.nn.functional.mse_loss(values_pred.squeeze(-1), batch_returns)\n\n                    # Total loss\n                    total_loss = actor_loss + config[\"training\"][\"critic_coeff\"] * critic_loss\n\n                    # Backward pass\n                    total_loss.backward()\n\n                    # Gradient clipping\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n\n                    optimizer.step()\n\n                # Logging\n                logger.debug(f\"Rank {rank}, Epoch {epoch + 1}, PPO Epoch {ppo_epoch + 1}, Loss: {total_loss.item():.6f}\")\n\n        # Reduce metrics across all processes\n        avg_reward = rewards.mean()\n        dist.all_reduce(avg_reward, op=dist.ReduceOp.SUM)\n        avg_reward = avg_reward.item() / world_size\n\n        logger.info(f\"Rank {rank}, Epoch {epoch + 1} - Average Reward: {avg_reward:.6f}\")\n\n        # Only rank 0 performs evaluation and checkpointing\n        if rank == 0:\n            # Evaluation\n            if (epoch + 1) % config['evaluation']['interval_eval'] == 0:\n                eval_model = model.module.actor if hasattr(model, \"module\") else model.actor\n                eval_metrics = evaluate_model(eval_model, env, config, device)\n            else:\n                eval_metrics = {}\n\n            # Monitoring\n            avg_metrics = {\n                \"actor_loss\": actor_loss.item(),\n                \"critic_loss\": critic_loss.item(),\n                \"entropy\": entropy.item()\n            }\n            monitor_training(epoch, rank, avg_metrics, eval_metrics, writer)\n\n            # Checkpointing\n            if avg_reward > best_reward:\n                best_reward = avg_reward\n                checkpoint_path = config[\"checkpoint\"][\"best_model_path\"]\n                torch.save(model.module.state_dict(), checkpoint_path)\n                logger.info(f\"New best model saved at {checkpoint_path} with reward {best_reward:.6f}\")\n\n        log_memory_usage(logger, phase=f\"Epoch {epoch + 1} End\")\n        log_gpu_memory(logger, phase=f\"Epoch {epoch + 1} End\")\n\n        # Synchronize processes at the end of each epoch\n        dist.barrier()\n\n        # Resource cleanup\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n        gc.collect()\n\n    if writer:\n        writer.close()\n    logger.info(f\"Training completed on Rank {rank}.\")",
    "./training/utils.py": "# File: training/utils.py\n\nimport os\nimport torch\nimport logging\nimport numpy as np\nimport random\nfrom typing import List, Tuple, Dict\nfrom functools import wraps\nimport psutil\nimport gc\nfrom pydantic import BaseModel, ValidationError\n\n# Initialize logger\nlogger = logging.getLogger(\"Utils\")\nlogger.setLevel(logging.DEBUG)\nif not logger.hasHandlers():\n    os.makedirs(\"logs\", exist_ok=True)\n    fh = logging.FileHandler(\"logs/utils.log\")\n    fh.setLevel(logging.DEBUG)\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n\ndef setup_logging(name, log_file=\"utils.log\"):\n    \"\"\"\n    Configure logging for utilities.\n    \"\"\"\n    return logger  # Logger is already set up\n\n\ndef setup_seed(seed):\n    \"\"\"\n    Set the random seed for reproducibility.\n    \"\"\"\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n    logger.info(f\"Random seed set to {seed}\")\n\n\ndef load_checkpoint(checkpoint_path, model, optimizer_actor, optimizer_critic, device):\n    logger.info(f\"Attempting to load checkpoint from {checkpoint_path}\")\n\n    if os.path.isfile(checkpoint_path):\n        checkpoint = torch.load(checkpoint_path, map_location=device)\n        model.load_state_dict(checkpoint[\"model_state_dict\"])\n        if optimizer_actor and optimizer_critic:\n            optimizer_actor.load_state_dict(checkpoint[\"optimizer_actor_state_dict\"])\n            optimizer_critic.load_state_dict(checkpoint[\"optimizer_critic_state_dict\"])\n        epoch = checkpoint[\"epoch\"]\n        best_reward = checkpoint.get(\"best_reward\", float(\"-inf\"))\n        logger.info(f\"Checkpoint loaded successfully from {checkpoint_path}, resuming at epoch {epoch}.\")\n        return epoch, best_reward\n    else:\n        logger.warning(f\"No checkpoint found at {checkpoint_path}.\")\n        return 0, float(\"-inf\")\n\n\ndef save_checkpoint(checkpoint_path, model, optimizer_actor, optimizer_critic, epoch, best_reward):\n    \"\"\"\n    Saves a checkpoint of the model and optimizer states.\n    \"\"\"\n    try:\n        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)\n        torch.save(\n            {\n                \"model_state_dict\": model.state_dict(),\n                \"optimizer_actor_state_dict\": optimizer_actor.state_dict() if optimizer_actor else None,\n                \"optimizer_critic_state_dict\": optimizer_critic.state_dict() if optimizer_critic else None,\n                \"epoch\": epoch,\n                \"best_reward\": best_reward,\n            },\n            checkpoint_path,\n        )\n        logger.info(f\"Checkpoint saved successfully at {checkpoint_path}.\")\n    except Exception as e:\n        logger.error(f\"Failed to save checkpoint: {e}\")\n\n\n# Define Nested Pydantic Models\nclass DataConfig(BaseModel):\n    raw_data_path: str\n    feature_columns: List[str]\n    n_steps: int\n    reward_scale: float\n    processed_data_path: str = \"data/processed/processed_data.csv\"\n    normalize_features: bool = True\n\n\nclass ModelConfig(BaseModel):\n    action_dim: int\n    feature_dim: int\n    ppo_hidden_dim: int\n    ppo_n_heads: int\n    ppo_n_layers: int\n    ppo_dropout: float\n    temperature: float = 1.0\n    layer_norm: bool = False\n\n\nclass TrainingConfig(BaseModel):\n    batch_size: int\n    epochs: int\n    actor_lr: float\n    critic_lr: float\n    clip_epsilon: float\n    gamma: float\n    gae_lambda: float\n    entropy_coeff: float\n    critic_coeff: float\n    ppo_epochs: int\n    weight_decay: float = 0.0\n    max_grad_norm: float = 0.5\n\n\nclass CheckpointConfig(BaseModel):\n    best_model_path: str\n    resume_checkpoint: str = \"\"\n\n\nclass EvaluationConfig(BaseModel):\n    interval_eval: int\n    episodes: int\n    baseline_reward: float\n\n\nclass LoggingConfig(BaseModel):\n    log_dir: str\n    use_tensorboard: bool\n\n\nclass DistributedConfig(BaseModel):\n    master_addr: str\n    master_port: str\n    world_size: int = 1\n\n\nclass CuriosityConfig(BaseModel):\n    use_curiosity: bool\n    learning_rate: float\n    hidden_dim_curiosity: int\n\n\nclass PreprocessingConfig(BaseModel):\n    enable: bool\n\n\nclass EnvironmentConfig(BaseModel):\n    max_steps: int\n\n\nclass ConfigModel(BaseModel):\n    random_seed: int\n    data: DataConfig\n    model: ModelConfig\n    training: TrainingConfig\n    checkpoint: CheckpointConfig\n    evaluation: EvaluationConfig\n    logging: LoggingConfig\n    distributed: DistributedConfig\n    curiosity: CuriosityConfig\n    preprocessing: PreprocessingConfig\n    environment: EnvironmentConfig\n\n\ndef validate_and_cast_config(config):\n    \"\"\"\n    Validate and cast the configuration dictionary to ensure all required keys\n    are present and of correct types.\n    \"\"\"\n    logger.info(\"Validating configuration parameters.\")\n\n    try:\n        validated_config = ConfigModel(**config)\n        logger.info(\"Configuration validated successfully.\")\n        return validated_config.dict()\n    except ValidationError as e:\n        logger.error(f\"Configuration validation failed: {e}\")\n        raise\n\n\ndef compute_advantages_and_returns(rewards, values, dones, gamma, gae_lambda, reward_scale=1.0):\n    \"\"\"\n    Compute Generalized Advantage Estimation (GAE) and returns.\n\n    Args:\n        rewards (Tensor): Tensor of rewards.\n        values (Tensor): Tensor of value predictions.\n        dones (Tensor): Tensor of done flags.\n        gamma (float): Discount factor for future rewards.\n        gae_lambda (float): GAE lambda for smoothing.\n        reward_scale (float, optional): Scaling factor for rewards.\n\n    Returns:\n        Tuple[Tensor, Tensor]: Advantages and returns tensors.\n    \"\"\"\n    rewards = rewards * reward_scale\n    advantages = torch.zeros_like(rewards)\n    returns = torch.zeros_like(rewards)\n    gae = 0\n    next_value = 0\n\n    for t in reversed(range(len(rewards))):\n        mask = 1.0 - dones[t]\n        delta = rewards[t] + gamma * next_value * mask - values[t]\n        gae = delta + gamma * gae_lambda * mask * gae\n        advantages[t] = gae\n        next_value = values[t]\n\n    returns = advantages + values\n    return advantages, returns\n\n\ndef log_memory_usage(logger, phase=\"\"):\n    \"\"\"\n    Logs the current CPU and memory usage.\n    \"\"\"\n    process = psutil.Process(os.getpid())\n    mem_info = process.memory_info()\n    mem_usage_mb = mem_info.rss / (1024 ** 2)  # Convert bytes to MB\n    cpu_percent = psutil.cpu_percent(interval=1)\n    logger.info(f\"Resource Usage [{phase}]: CPU: {cpu_percent}% | Memory: {mem_usage_mb:.2f} MB\")\n\n\ndef log_gpu_memory(logger, phase=\"\"):\n    \"\"\"\n    Logs the current GPU memory usage.\n    \"\"\"\n    if torch.cuda.is_available():\n        allocated = torch.cuda.memory_allocated() / (1024 ** 2)  # in MB\n        reserved = torch.cuda.memory_reserved() / (1024 ** 2)     # in MB\n        logger.info(f\"GPU Memory Usage [{phase}]: Allocated: {allocated:.2f} MB | Reserved: {reserved:.2f} MB\")\n    else:\n        logger.info(f\"GPU Memory Usage [{phase}]: GPU not available.\")",
    "./training/__init__.py": "",
    "./training/evaluation.py": "# File: training/evaluation.py\n\nimport torch\nimport numpy as np\nimport logging\nfrom collections import Counter\nfrom training.utils import log_memory_usage, log_gpu_memory\nimport gc\n\n\ndef evaluate_model(actor, environment, config, device):\n    \"\"\"\n    Evaluate the model with comprehensive financial metrics.\n\n    Args:\n        actor: Actor network for policy evaluation.\n        environment: Trading environment.\n        config: Full configuration dictionary.\n        device: Computation device (e.g., 'cpu' or 'cuda').\n\n    Returns:\n        Dictionary containing evaluation metrics.\n    \"\"\"\n    logger = logging.getLogger(\"Evaluation\")\n    logger.info(\"Starting model evaluation.\")\n    log_memory_usage(logger, phase=\"Before Evaluation\")\n    log_gpu_memory(logger, phase=\"Before Evaluation\")\n\n    actor.eval()  # Set actor to evaluation mode\n    total_rewards = []\n    sharpe_ratios = []\n    max_drawdowns = []\n    trade_counts = []\n    win_rates = []\n\n    num_episodes = config[\"evaluation\"][\"episodes\"]\n    gamma = config[\"training\"].get(\"gamma\", 0.99)  # Discount factor used in returns\n    reward_scale = config['data'].get('reward_scale', 1.0)\n\n    for episode in range(num_episodes):\n        state = environment.reset()\n        episode_reward = 0.0\n        cumulative_returns = 0.0\n        peak = 0.0\n        done = False\n        steps = 0\n        max_steps_per_episode = config[\"environment\"][\"max_steps\"]\n        action_counts = Counter({\"Sell\": 0, \"Hold\": 0, \"Buy\": 0})\n        episode_returns = []\n\n        while not done and steps < max_steps_per_episode:\n            state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(device)\n            with torch.no_grad():\n                logits = actor(state_tensor)\n                action = torch.argmax(logits, dim=-1).item()\n\n            step_result = environment.step(action)\n            next_state, reward, done, info = step_result\n\n            scaled_reward = reward * reward_scale\n            cumulative_returns += scaled_reward\n            episode_reward += scaled_reward\n            peak = max(peak, cumulative_returns)\n            episode_returns.append(scaled_reward)\n\n            # Simple action mapping for logging\n            action_label = [\"Sell\", \"Hold\", \"Buy\"][action]\n            action_counts[action_label] += 1\n\n            steps += 1\n            state = next_state\n\n        total_rewards.append(episode_reward)\n\n        # Compute Sharpe ratio\n        avg_return = np.mean(episode_returns) if len(episode_returns) > 0 else 0\n        std_dev_return = np.std(episode_returns) + 1e-8\n        # Annualized Sharpe ratio (assuming 252 trading days)\n        sharpe_ratio = (avg_return / std_dev_return) * np.sqrt(252)\n        sharpe_ratios.append(sharpe_ratio)\n\n        # Compute max drawdown\n        max_drawdown = peak - cumulative_returns\n        max_drawdowns.append(max_drawdown)\n\n        # Compute win rate\n        wins = sum(1 for r in episode_returns if r > 0)\n        total_trades = len(episode_returns)\n        win_rate = wins / total_trades if total_trades > 0 else 0\n        win_rates.append(win_rate)\n\n        trade_counts.append(total_trades)\n        logger.debug(\n            f\"Episode {episode + 1}: Reward={episode_reward:.6f}, Trades={total_trades}, \"\n            f\"Sharpe Ratio={sharpe_ratio:.6f}, Max Drawdown={max_drawdown:.6f}, Win Rate={win_rate:.2%}\"\n        )\n        logger.info(f\"Action counts: {action_counts}\")\n\n        # Resource cleanup\n        del state_tensor\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n        gc.collect()\n\n    # Aggregate metrics\n    avg_reward = np.mean(total_rewards)\n    avg_sharpe_ratio = np.mean(sharpe_ratios)\n    avg_max_drawdown = np.mean(max_drawdowns)\n    avg_trade_count = np.mean(trade_counts)\n    avg_win_rate = np.mean(win_rates)\n\n    metrics = {\n        \"average_reward\": avg_reward,\n        \"sharpe_ratio\": avg_sharpe_ratio,\n        \"max_drawdown\": avg_max_drawdown,\n        \"average_trade_count\": avg_trade_count,\n        \"average_win_rate\": avg_win_rate,\n    }\n\n    logger.info(f\"Evaluation Metrics: {metrics}\")\n    log_memory_usage(logger, phase=\"After Evaluation\")\n    log_gpu_memory(logger, phase=\"After Evaluation\")\n    return metrics",
    "./config/config.yaml": "# config/config.yaml\n\n# =====================\n# Random Seed for Reproducibility\n# =====================\nrandom_seed: 42\n\n# =====================\n# Logging Configuration\n# =====================\nlogging:\n  log_dir: \"./logs\"          # Directory to store log files\n  use_tensorboard: true      # Whether to use TensorBoard\n\n# =====================\n# Distributed Training Configuration\n# =====================\ndistributed:\n  master_addr: \"localhost\"   # Master node address\n  master_port: \"29500\"       # Master node port\n  world_size: 2              # Number of GPUs to use for training (Note: Ensure this is used correctly in your code)\n\n# =====================\n# Data Configuration\n# =====================\ndata:\n  raw_data_path: \"./data/dataset.csv\"                     # Path to the raw data CSV file\n  feature_columns:\n    - open\n    - high\n    - low\n    - close\n    - volume\n    - elapsed_time\n    - rsi\n    - adx\n    - roc\n    - stochastic_k\n    - stochastic_d\n    - macd\n    - macd_signal\n    - macd_hist\n    - mfi\n    - bb_upper\n    - bb_middle\n    - bb_lower\n    - kc_upper\n    - kc_lower\n    - dc_upper\n    - dc_lower\n    - atr\n    - pivot\n    - fibonacci_23\n    - fibonacci_38\n    - fibonacci_50\n    - fibonacci_61\n    - vpt\n    - cmf\n    - obv\n    - supertrend\n    # Add other feature columns as needed\n  n_steps: 20                                             # Number of steps to consider in each sample\n  reward_scale: 1.0                                   # Scaling factor for rewards\n  pt_data_dir: \"./data/processed/\"  \n  num_splits: 10\n  normalize_features: true                                # Whether to normalize features\n\n# =====================\n# Model Configuration\n# =====================\nmodel:\n  feature_dim: 32              # Number of input features\n  ppo_hidden_dim: 384         # Hidden size of the model layers\n  action_dim: 3                # Number of possible actions (e.g., Sell=0, Hold=1, Buy=2)\n  ppo_n_heads: 6              # Number of attention heads in the Transformer\n  ppo_n_layers: 3              # Number of Transformer encoder layers\n  ppo_dropout: 0.1             # Dropout rate for regularization\n  temperature: 1.0             # Temperature parameter for action selection\n  layer_norm: true             # Whether to apply layer normalization\n\n# =====================\n# Training Configuration\n# =====================\ntraining:\n  epochs: 100\n  batch_size: 512\n  actor_lr: 1e-4\n  critic_lr: 1e-5\n  ppo_epochs: 4\n  clip_epsilon: 0.2\n  entropy_coeff: 0.1\n  critic_coeff: 0.5\n  gamma: 0.99\n  gae_lambda: 0.95\n  weight_decay: 0.0\n  max_grad_norm: 0.2\n\n# =====================\n# Checkpointing Configuration\n# =====================\ncheckpoint:\n  best_model_path: \"./checkpoints/best_model.pth\"  # Path to save the best model\n  resume_checkpoint: \"\"                            # Path to resume training from (leave empty if not resuming)\n\n# =====================\n# Evaluation Configuration\n# =====================\nevaluation:\n  interval_eval: 5      # Frequency of evaluation (in epochs)\n  episodes: 20\n  baseline_reward: 0.0\n  gamma: 0.99            # Discount factor used during evaluation\n\n# =====================\n# Environment Configuration\n# =====================\nenvironment:\n  max_steps: 10000        # Maximum steps per episode\n\n# =====================\n# Curiosity-Driven Exploration Configuration (Optional)\n# =====================\ncuriosity:\n  use_curiosity: false\n  learning_rate: 0.001\n  hidden_dim_curiosity: 256\n\n# =====================\n# Preprocessing Configuration\n# =====================\npreprocessing:\n  enable: true",
    "./models/models.py": "# File: models/models.py\n\nimport os\nimport torch\nimport torch.nn as nn\nfrom models.transformer import TransformerEncoder\nimport logging\nimport torch.nn.init as init\n\n# Initialize logger for Models\nlogger = logging.getLogger(\"Models\")\nlogger.setLevel(logging.DEBUG)\nif not logger.hasHandlers():\n    os.makedirs(\"logs\", exist_ok=True)\n    fh = logging.FileHandler(\"logs/models.log\")\n    fh.setLevel(logging.DEBUG)\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n\ndef initialize_weights(m):\n    \"\"\"\n    Initialize weights of the model.\n    Use Xavier Uniform for Tanh and He initialization for LeakyReLU.\n    \"\"\"\n    if isinstance(m, nn.Linear):\n        init.kaiming_uniform_(m.weight, nonlinearity='leaky_relu')\n        if m.bias is not None:\n            init.constant_(m.bias, 0)\n    elif isinstance(m, nn.LayerNorm):\n        init.constant_(m.weight, 1)\n        init.constant_(m.bias, 0)\n\n    # Check for NaN/Inf after initialization (optional)\n    if hasattr(m, 'weight') and (torch.isnan(m.weight).any() or torch.isinf(m.weight).any()):\n        logger.error(f\"NaN/Inf detected in weights of {type(m).__name__}\")\n    if hasattr(m, 'bias') and m.bias is not None and (torch.isnan(m.bias).any() or torch.isinf(m.bias).any()):\n        logger.error(f\"NaN/Inf detected in bias of {type(m).__name__}\")\n\n\nclass Actor(nn.Module):\n    def __init__(self, feature_dim, hidden_dim, action_dim, n_heads, n_layers, temperature=1.0, layer_norm=False):\n        super().__init__()\n        self.encoder = TransformerEncoder(feature_dim, hidden_dim, n_heads, n_layers, dropout=0.1)\n        self.layer_norm_flag = layer_norm\n        if self.layer_norm_flag:\n            self.layer_norm = nn.LayerNorm(hidden_dim)\n\n        self.fc = nn.Linear(hidden_dim, action_dim)  # Output logits for actions\n        self.temperature = temperature\n\n    def forward(self, x):\n        x = self.encoder(x)\n        x = x[:, -1]  # Get the output from the last time step\n        if self.layer_norm_flag:\n            x = self.layer_norm(x)\n        logits = self.fc(x)\n        if torch.isnan(logits).any():\n            logger.error(\"Logits contain NaNs in Actor forward pass\")\n            raise ValueError(\"Logits contain NaNs in Actor\")\n        if self.temperature != 1.0:\n            logits = logits / self.temperature  # Apply temperature scaling\n\n        return logits\n\n\nclass Critic(nn.Module):\n    def __init__(self, feature_dim, hidden_dim, n_heads, n_layers, layer_norm=False):\n        super().__init__()\n        self.encoder = TransformerEncoder(feature_dim, hidden_dim, n_heads, n_layers, dropout=0.1)\n        self.layer_norm_flag = layer_norm\n        if self.layer_norm_flag:\n            self.layer_norm = nn.LayerNorm(hidden_dim)\n\n        self.fc = nn.Sequential(\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.LeakyReLU(negative_slope=0.01),\n            nn.Linear(hidden_dim, 1)\n            # Removed Tanh activation to allow value predictions beyond [-1,1]\n        )\n\n    def forward(self, x):\n        x = self.encoder(x)\n        x = x[:, -1]  # Get the output from the last time step\n        if self.layer_norm_flag:\n            x = self.layer_norm(x)\n        values = self.fc(x)\n        if torch.isnan(values).any():\n            logger.error(\"Values contain NaNs in Critic forward pass\")\n            raise ValueError(\"Values contain NaNs in Critic\")\n        return values\n\n\nclass PPOAgent(nn.Module):\n    def __init__(self, feature_dim, hidden_dim, action_dim, n_heads, n_layers,\n                 dropout=0.1, temperature=1.0, layer_norm=False):\n        super().__init__()\n        self.actor = Actor(feature_dim, hidden_dim, action_dim, n_heads, n_layers, temperature, layer_norm)\n        self.critic = Critic(feature_dim, hidden_dim, n_heads, n_layers, layer_norm)\n        self.apply(initialize_weights)\n        logger.info(\"PPOAgent weights initialized.\")\n\n    def forward(self, x):\n        logits = self.actor(x)\n        value = self.critic(x)\n        return logits, value",
    "./models/transformer.py": "# File: models/transformer.py\n\nimport os\nimport torch\nimport torch.nn as nn\nimport math\nimport logging\n\nlogger = logging.getLogger(\"DebugTransformer\")\nlogger.setLevel(logging.DEBUG)\nif not logger.hasHandlers():\n    os.makedirs(\"logs\", exist_ok=True)\n    fh = logging.FileHandler(\"logs/transformer.log\")\n    fh.setLevel(logging.DEBUG)\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=5000):\n        super().__init__()\n        self.encoding = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len).unsqueeze(1).float()\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))\n        self.encoding[:, 0::2] = torch.sin(position * div_term)\n        self.encoding[:, 1::2] = torch.cos(position * div_term)\n        self.encoding = self.encoding.unsqueeze(0)\n\n    def forward(self, x):\n        return x + self.encoding[:, :x.size(1), :].to(x.device)\n\n\nclass TransformerEncoder(nn.Module):\n    def __init__(self, input_dim, hidden_dim, n_heads, n_layers, dropout=0.1):\n        super().__init__()\n        self.embedding = nn.Linear(input_dim, hidden_dim)\n        self.pos_enc = PositionalEncoding(hidden_dim)\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=hidden_dim, nhead=n_heads, dropout=dropout, batch_first=True\n        )\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x):\n        x = self.embedding(x)\n        logger.debug(f\"TransformerEncoder embed stats => min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}\")\n        x = self.pos_enc(x)\n        logger.debug(f\"After pos_enc => min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}\")\n        x = self.dropout(x)\n        x = self.transformer_encoder(x)\n        logger.debug(f\"After transformer => min={x.min():.4f}, max={x.max():.4f}, mean={x.mean():.4f}\")\n        return x",
    "./data_processing/preprocess.py": "# File: data_processing/preprocess.py\n\nimport os\nimport logging\nimport pandas as pd\nimport numpy as np\nimport joblib\nimport torch\nfrom sklearn.preprocessing import MinMaxScaler\nfrom typing import Dict\n\nlogger = logging.getLogger(\"Preprocessing\")\nlogger.setLevel(logging.DEBUG)\nif not logger.hasHandlers():\n    os.makedirs(\"logs\", exist_ok=True)\n    fh = logging.FileHandler(\"logs/preprocessing.log\")\n    fh.setLevel(logging.DEBUG)\n    ch = logging.StreamHandler()\n    ch.setLevel(logging.INFO)\n    formatter = logging.Formatter(\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\")\n    fh.setFormatter(formatter)\n    ch.setFormatter(formatter)\n    logger.addHandler(fh)\n    logger.addHandler(ch)\n\n\ndef export_and_save_preprocessed_data(config: Dict, logger: logging.Logger) -> None:\n    \"\"\"\n    Preprocess the raw data, split into multiple PT files.\n    \"\"\"\n    raw_data_path = config['data']['raw_data_path']\n    pt_data_dir = config['data'].get('pt_data_dir', 'data/processed/')\n    pt_data_prefix = 'processed_data_part'\n    os.makedirs(pt_data_dir, exist_ok=True)\n\n    # Check if already preprocessed\n    if any(fname.startswith(pt_data_prefix) for fname in os.listdir(pt_data_dir)):\n        logger.info(f\"Processed data files already exist in '{pt_data_dir}'. Skipping preprocessing.\")\n        return\n\n    try:\n        logger.info(f\"Reading raw data from: {raw_data_path}\")\n        df = pd.read_csv(raw_data_path)\n        logger.info(f\"Initial data shape: {df.shape}\")\n\n        # Replace infinite values with NaN\n        df.replace([np.inf, -np.inf], np.nan, inplace=True)\n\n        # Drop rows with NaN or infinite values\n        initial_length = len(df)\n        df.dropna(inplace=True)\n        cleaned_length = len(df)\n        logger.info(f\"Removed {initial_length - cleaned_length} rows containing NaN or infinite values.\")\n\n        # Check if dataframe is empty after cleaning\n        if df.empty:\n            logger.error(\"Preprocessed dataframe is empty after removing NaNs and infinite values.\")\n            raise ValueError(\"No data left after preprocessing.\")\n\n        # Normalize features if required\n        feat_cols = config['data']['feature_columns']\n\n        # Ensure 'close' is in feature columns\n        if 'close' not in feat_cols:\n            feat_cols.append('close')\n\n        normalize_features = config['data'].get('normalize_features', True)\n        if normalize_features:\n            logger.info(\"Normalizing feature columns.\")\n            df = normalize_feature_columns(df, feat_cols, logger)\n\n        # Split data into parts and save as PT files\n        num_splits = config['data'].get('num_splits', 10)\n        data_splits = np.array_split(df[feat_cols].values, num_splits)\n\n        for idx, split_data in enumerate(data_splits):\n            data_tensor = torch.tensor(split_data, dtype=torch.float32)\n            pt_file_path = os.path.join(pt_data_dir, f\"{pt_data_prefix}_{idx}.pt\")\n            torch.save(data_tensor, pt_file_path)\n            logger.info(f\"Saved split {idx} to {pt_file_path}\")\n\n    except Exception as e:\n        logger.error(f\"Failed to preprocess data: {e}\", exc_info=True)\n        raise\n\n\ndef normalize_feature_columns(df: pd.DataFrame, feat_cols: list, logger: logging.Logger) -> pd.DataFrame:\n    \"\"\"\n    Normalizes the feature columns using MinMaxScaler, excluding 'close'.\n    \"\"\"\n    scaler_path = os.path.join('scalers', 'minmax_scaler.joblib')\n    os.makedirs('scalers', exist_ok=True)\n\n    # Exclude 'close' from normalization\n    features_to_normalize = [col for col in feat_cols if col != 'close']\n\n    if os.path.exists(scaler_path):\n        logger.info(f\"Loading scaler from {scaler_path}\")\n        scaler = joblib.load(scaler_path)\n        # Check if features match\n        if hasattr(scaler, 'feature_names_in_'):\n            scaler_features = scaler.feature_names_in_\n            if not np.array_equal(scaler_features, features_to_normalize):\n                logger.warning(\"Scaler feature names do not match features to normalize. Re-fitting the scaler.\")\n                scaler = MinMaxScaler(feature_range=(0, 1))\n                scaler.fit(df[features_to_normalize])\n                joblib.dump(scaler, scaler_path)\n                logger.info(f\"Scaler re-fitted and saved to {scaler_path}\")\n        else:\n            logger.warning(\"Scaler does not have 'feature_names_in_'. Re-fitting the scaler.\")\n            scaler = MinMaxScaler(feature_range=(0, 1))\n            scaler.fit(df[features_to_normalize])\n            joblib.dump(scaler, scaler_path)\n            logger.info(f\"Scaler re-fitted and saved to {scaler_path}\")\n\n        df[features_to_normalize] = scaler.transform(df[features_to_normalize])\n    else:\n        logger.info(f\"No scaler found at {scaler_path}. Fitting a new scaler.\")\n        scaler = MinMaxScaler(feature_range=(0, 1))\n        scaler.fit(df[features_to_normalize])\n        df[features_to_normalize] = scaler.transform(df[features_to_normalize])\n        joblib.dump(scaler, scaler_path)\n        logger.info(f\"Scaler saved to {scaler_path}\")\n\n    return df"
}