ipynb_baseline_code.txt

# %% [markdown]
# # A3C for Pong (ALE/Pong-v5) with Gymnasium & PyTorch on Colab
#
# 이 노트북은 Google Colab 환경에서 Gymnasium의 `ALE/Pong-v5` 환경을 사용하여 Asynchronous Advantage Actor-Critic (A3C) 알고리즘으로 강화학습 에이전트를 훈련합니다.
# 멀티프로세싱을 사용하여 여러 워커가 병렬로 학습을 진행하며, 주기적으로 학습 과정 영상과 모델을 저장합니다.

# %% [markdown]
# ## 0. 패키지 설치 (Colab 환경)
# Colab 환경에서 코드를 실행하기 위해 필요한 패키지들을 설치합니다. 로컬 환경에서는 이미 설치되어 있거나 `requirements.txt`를 통해 설치할 수 있습니다.

# %%
# Colab에서 실행 시 다음 셀의 주석을 해제하고 실행하세요.
# !pip install gymnasium[atari,accept-rom-license] ale-py pyvirtualdisplay imageio opencv-python tqdm matplotlib torch --quiet

# %% [markdown]
# ## 1. 라이브러리 임포트
# 필요한 모든 라이브러리를 임포트합니다.

# %%
import os
import random
import numpy as np
import gymnasium as gym
from gymnasium.wrappers.atari_preprocessing import AtariPreprocessing
from gymnasium.wrappers import FrameStackObservation # Gymnasium v0.26+ (실제로는 v1.0.0a1+)
import ale_py  # Atari ROM 라이선스 동의 후 ALE 환경 사용에 필요

from tqdm import tqdm
import pickle
import imageio
import time
import json
from datetime import datetime
import glob
from pyvirtualdisplay import Display # Colab 또는 headless 서버에서 렌더링을 위함

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.multiprocessing as mp
from collections import deque
import cv2 # imageio가 내부적으로 사용하거나, 프레임 처리에 필요할 수 있음

# %% [markdown]
# ## 2. 설정값 및 하이퍼파라미터
# 학습에 사용될 주요 설정값과 하이퍼파라미터를 정의합니다.

# %%
# --- 설정값 ---
MAX_GLOBAL_STEPS = 200000  # 총 학습 스텝 수
RECORD_INTERVAL_GLOBAL_STEPS = 50000  # 영상 녹화 간격 (글로벌 스텝 기준)
N_WORKERS = 4  # 워커 수 (Colab의 CPU 코어 수에 맞춰 조절 가능, 보통 2 또는 4)
ENV_NAME = 'ALE/Pong-v5'  # 환경 이름 (Gymnasium Atari 환경)
LEARNING_RATE = 1e-4       # 학습률
GAMMA = 0.99               # 할인율
ENTROPY_BETA = 0.01        # 엔트로피 정규화 계수
N_STEPS_UPDATE = 20        # n-스텝 업데이트 주기

# %% [markdown]
# ## 3. 출력 디렉토리 생성 및 디스플레이 설정
# 학습 결과(모델, 로그, 비디오)를 저장할 디렉토리를 생성하고, Colab 환경을 위한 가상 디스플레이를 설정합니다.

# %%
# 출력 디렉토리 생성
OUTPUT_DIR = "output_a3c_pong_colab"
os.makedirs(OUTPUT_DIR, exist_ok=True)
MODELS_DIR = os.path.join(OUTPUT_DIR, "models")
os.makedirs(MODELS_DIR, exist_ok=True)
LOGS_DIR = os.path.join(OUTPUT_DIR, "logs") # 로그는 현재 코드에서 명시적으로 저장하진 않음
os.makedirs(LOGS_DIR, exist_ok=True)
VIDEOS_DIR = os.path.join(OUTPUT_DIR, "videos")
os.makedirs(VIDEOS_DIR, exist_ok=True)

print(f"출력 디렉토리: {os.path.abspath(OUTPUT_DIR)}")

# %%
# 가상 디스플레이 설정 (Colab에서 GUI 없이 렌더링하기 위함)
try:
    display = Display(visible=0, size=(400, 300)) # Pong 화면 크기에 맞춰 조정 가능
    display.start()
    print("가상 디스플레이가 시작되었습니다.")
except Exception as e:
    print(f"가상 디스플레이를 시작할 수 없습니다: {e}. 영상 저장이 제대로 안될 수 있습니다.")
    # Colab에서는 display.stop()을 명시적으로 호출하지 않아도 세션 종료 시 정리될 수 있음

# %% [markdown]
# ## 4. 헬퍼 함수 정의
# 학습 과정에서 사용될 유틸리티 함수들을 정의합니다.

# %% [markdown]
# ### 4.1 환경 생성 함수
# Atari 환경을 생성하고 필요한 전처리 래퍼(Wrapper)를 적용합니다.
# `AtariPreprocessing`은 프레임 스킵, 그레이스케일 변환, 관찰 크기 조정, 생명 손실 시 에피소드 종료 등을 처리합니다.
# `FrameStackObservation`은 여러 프레임을 쌓아 하나의 관찰로 만듭니다.

# %%
def create_env(render_mode_param=None):
    """Atari 환경 생성 및 전처리 래퍼 적용 함수"""
    # 원본 환경의 프레임 스킵은 1로 설정 (AtariPreprocessing에서 자체 프레임 스킵 사용)
    env = gym.make(ENV_NAME, render_mode=render_mode_param, frameskip=1)
    
    # AtariPreprocessing 래퍼 적용
    env = AtariPreprocessing(env,
                             frame_skip=4,            # 4 프레임마다 하나의 관찰 생성
                             grayscale_obs=True,      # 관찰을 그레이스케일로 변환
                             scale_obs=False,         # 관찰 값을 0-1로 스케일링 (True로 하면 신경망 입력에 유리할 수 있으나, 여기선 False 후 preprocess_frame에서 처리 가능성)
                                                      # -> ActorCriticNetwork는 0-255 범위의 이미지를 받도록 되어있음 (나중에 /255.0 처리)
                                                      # scale_obs=True로 하면 0-1 범위가 됨. 네트워크 입력 전에 /255 안해도 됨.
                                                      # 여기서는 False로 두고, get_stacked_frames에서 /255.0 처리
                             terminal_on_life_loss=True) # 생명(life)을 잃으면 에피소드 종료
    
    # 프레임 스택 래퍼 적용
    env = FrameStackObservation(env, stack_size=4) # 4개의 프레임을 쌓음
    return env

# %% [markdown]
# ### 4.2 비디오 저장 함수

# %%
def save_video(frames_list, video_filename, fps=15):
    """수집된 프레임 리스트를 사용하여 비디오 파일을 저장합니다."""
    try:
        imageio.mimsave(
            video_filename,
            frames_list,
            fps=fps,      # AtariPreprocessing의 frame_skip=4를 고려하여 실제 게임 속도와 유사하게 설정
            quality=8,    # 비디오 품질 (0-10, 높을수록 좋음)
            macro_block_size=1 # 더 부드러운 영상 위한 설정 (코덱에 따라 다름)
        )
        print(f"비디오가 저장되었습니다: {video_filename}")
    except Exception as e:
        print(f"비디오 저장 중 오류 발생 {video_filename}: {e}")

# %% [markdown]
# ## 5. Actor-Critic 신경망 정의
# A3C 알고리즘의 핵심이 되는 Actor(정책망)와 Critic(가치망) 역할을 동시에 수행하는 CNN 기반 신경망을 정의합니다.

# %%
class ActorCriticNetwork(nn.Module):
    def __init__(self, num_input_channels, num_actions):
        super(ActorCriticNetwork, self).__init__()

        # 공유 특징 추출 레이어 (Convolutional Layers)
        # AtariPreprocessing 및 FrameStackObservation을 거친 입력은 (4, 84, 84) 형태가 됨
        self.conv1 = nn.Conv2d(num_input_channels, 32, kernel_size=8, stride=4) # (N, 4, 84, 84) -> (N, 32, 20, 20)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)                   # (N, 32, 20, 20) -> (N, 64, 9, 9)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)                   # (N, 64, 9, 9) -> (N, 64, 7, 7)

        # 완전 연결 레이어를 위한 Flatten 후 크기 계산
        # 입력 이미지 (84,84) 가정
        dummy_input = torch.zeros(1, num_input_channels, 84, 84)
        self.conv_output_size = self._get_conv_output_size(dummy_input)

        # 공유 FC 레이어
        self.fc_shared = nn.Linear(self.conv_output_size, 512)

        # 정책 헤드 (Actor) - 행동 로짓(logits) 출력
        self.policy_head = nn.Linear(512, num_actions)

        # 가치 헤드 (Critic) - 상태 가치 출력
        self.value_head = nn.Linear(512, 1)

    def _get_conv_output_size(self, x_shape_tensor):
        x = F.relu(self.conv1(x_shape_tensor))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        return int(np.prod(x.size()[1:])) # 배치 차원(0) 제외하고 나머지 차원 곱

    def forward(self, x_state_tensor):
        # 입력 텐서는 (N, C, H, W) 형태여야 하고, 값의 범위는 0.0 ~ 1.0 이어야 함 (AtariPreprocessing(scale_obs=True) 또는 수동 스케일링)
        # 현재 get_stacked_frames에서 / 255.0 처리
        
        # 공유 특징 추출 네트워크 (CNN)
        x = F.relu(self.conv1(x_state_tensor))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = x.view(x.size(0), -1)  # 평탄화 (Flatten)

        # 공유 완전 연결 레이어
        shared_features = F.relu(self.fc_shared(x))

        # 정책 헤드 (Actor): 행동 확률 로짓(logits) 반환
        action_logits = self.policy_head(shared_features)
        # Softmax는 손실 함수(CrossEntropyLoss)나 확률 분포 샘플링(Categorical) 시 내부적으로 처리됨
        # 여기서는 확률 분포를 직접 반환 (이전 코드와 일관성)
        action_probs = F.softmax(action_logits, dim=1)

        # 가치 헤드 (Critic): 상태 가치 반환
        state_value = self.value_head(shared_features)

        return action_probs, state_value

# %% [markdown]
# ## 6. A3C 워커(에이전트) 클래스 정의
# 각 워커는 독립적인 프로세스로 실행되며, 환경과 상호작용하고, 로컬 네트워크를 통해 그라디언트를 계산한 후 글로벌 네트워크를 업데이트합니다.

# %%
class A3CWorker(mp.Process):
    def __init__(self, shared_global_network, global_optimizer_instance, global_episode_counter,
                 worker_process_id, video_save_dir, record_video_interval,
                 n_steps_for_update=N_STEPS_UPDATE, discount_factor=GAMMA, entropy_regularization_beta=ENTROPY_BETA,
                 max_global_training_steps=MAX_GLOBAL_STEPS):
        super(A3CWorker, self).__init__()

        self.worker_id = worker_process_id
        self.global_network = shared_global_network
        self.global_optimizer = global_optimizer_instance
        self.global_counter = global_episode_counter # 스텝 카운터로 변경
        self.max_global_steps = max_global_training_steps

        # 영상 기록 관련
        self.video_dir = video_save_dir
        self.record_interval_global_steps = record_video_interval
        self.last_recorded_milestone = 0 # 마지막으로 녹화 시작한 글로벌 스텝 구간

        # 하이퍼파라미터
        self.n_steps = n_steps_for_update
        self.gamma = discount_factor
        self.entropy_beta = entropy_regularization_beta

        # 환경 생성 (워커별 독립 환경)
        # 영상 녹화는 worker 0만 담당하므로, render_mode를 다르게 설정
        if self.worker_id == 0:
            self.env = create_env(render_mode_param='rgb_array') # worker 0은 렌더링 가능한 환경
            print(f"워커 {self.worker_id}: render_mode='rgb_array'로 환경 생성")
        else:
            self.env = create_env() # 다른 워커는 렌더링 불필요
        
        # 네트워크 입력 채널 수 (FrameStackObservation의 stack_size와 동일)
        self.num_input_channels = self.env.observation_space.shape[0] # (C, H, W) 중 C
        num_actions = self.env.action_space.n

        # 로컬 네트워크 (글로벌 네트워크와 동일 구조)
        self.local_network = ActorCriticNetwork(self.num_input_channels, num_actions)
        self.total_episodes_by_worker = 0 # 해당 워커가 완료한 에피소드 수


    def convert_obs_to_tensor(self, stacked_observation_from_env):
        """환경에서 받은 스택된 관찰(LazyFrames)을 PyTorch 텐서로 변환하고 스케일링합니다."""
        # LazyFrames를 NumPy 배열로 변환하고, 타입을 float32로, 값의 범위를 0.0 ~ 1.0으로 스케일링
        np_array_obs = np.array(stacked_observation_from_env, dtype=np.float32) / 255.0
        # (1, C, H, W) 형태로 배치 차원 추가
        return torch.FloatTensor(np_array_obs).unsqueeze(0)

    def sync_local_network_with_global(self):
        """글로벌 네트워크의 파라미터를 로컬 네트워크로 복사합니다."""
        self.local_network.load_state_dict(self.global_network.state_dict())

    def calculate_n_step_loss(self, rewards_buffer, values_buffer, log_probs_buffer, entropies_buffer, is_episode_done):
        """n-스텝 후 손실(loss)을 계산합니다."""
        # 부트스트랩을 위한 R 값 초기화
        R_bootstrap = torch.zeros(1, 1) # (1,1) 형태의 텐서
        # 에피소드가 끝나지 않았고, value 버퍼가 비어있지 않다면, 마지막 상태의 가치로 R 설정
        if not is_episode_done and values_buffer:
            R_bootstrap = values_buffer[-1].detach() # 마지막 value를 사용 (s_t+n의 가치로 사용)

        policy_gradient_loss = 0
        value_function_loss = 0
        
        # 역순으로 n-스텝 동안의 보상, 어드밴티지, 손실 계산
        # R_bootstrap이 여기서 n-step 후의 가치 V(s_t+n) 역할을 함
        current_R = R_bootstrap
        for i in reversed(range(len(rewards_buffer))):
            current_R = rewards_buffer[i] + self.gamma * current_R # n-step return G_t:t+n
            advantage = current_R - values_buffer[i] # A_t = G_t:t+n - V(s_t)
            
            # 가치 손실 (Critic loss)
            value_function_loss += 0.5 * advantage.pow(2) # (A_t)^2 / 2
            
            # 정책 손실 (Actor loss): -log(pi(a_t|s_t)) * A_t - beta * H(pi(s_t))
            policy_gradient_loss += -(log_probs_buffer[i] * advantage.detach() + self.entropy_beta * entropies_buffer[i])
        
        # 총 손실 (하나의 스칼라 값으로 합산)
        # value_function_loss는 이미 스칼라들의 합이므로 .sum() 불필요. 각 스텝의 loss를 더했으므로.
        # 단, 여러 스텝에 대한 평균을 내고 싶다면 len(rewards_buffer)로 나눠야함. 여기서는 합산.
        if isinstance(value_function_loss, torch.Tensor) and value_function_loss.numel() > 1:
             value_function_loss = value_function_loss.sum()

        total_loss = policy_gradient_loss + value_function_loss
        return total_loss

    def run_worker_process(self):
        """워커의 메인 실행 루프입니다."""
        self.total_episodes_by_worker = 0
        while self.global_counter.value < self.max_global_steps:
            # 1. 글로벌 네트워크와 로컬 네트워크 동기화
            self.sync_local_network_with_global()

            # 2. 경험 저장을 위한 로컬 버퍼 초기화
            log_probs_buffer = []
            values_buffer = []
            rewards_buffer = []
            entropies_buffer = []

            # 3. 에피소드 시작 또는 이어서 진행
            current_observation_stacked, _ = self.env.reset() # (LazyFrames, info)
            current_state_tensor = self.convert_obs_to_tensor(current_observation_stacked)

            is_episode_done = False
            current_episode_reward = 0
            
            # --- 영상 녹화 설정 ---
            trigger_video_recording = False
            frames_for_current_video = []

            if self.worker_id == 0: # 워커 0만 영상 녹화
                # 현재 글로벌 스텝이 마지막 녹화 마일스톤 + 인터벌을 넘었는지 확인
                if self.global_counter.value >= self.last_recorded_milestone + self.record_interval_global_steps:
                    trigger_video_recording = True
                    # 다음 녹화 마일스톤 업데이트 (정확한 간격을 위해)
                    self.last_recorded_milestone = ((self.global_counter.value // self.record_interval_global_steps) + 1) * self.record_interval_global_steps
                    print(f"워커 0: 글로벌 스텝 ~{self.global_counter.value} 부근에서 에피소드 녹화 시작 (다음 마일스톤: {self.last_recorded_milestone})")
            
            if trigger_video_recording:
                # 환경의 render() 메소드는 현재 상태의 RGB 배열을 반환
                rendered_rgb_frame = self.env.render()
                if rendered_rgb_frame is not None:
                    frames_for_current_video.append(rendered_rgb_frame)
            # --- 영상 녹화 설정 끝 ---

            # 4. n-스텝 동안 환경과 상호작용 또는 에피소드 종료까지
            for step_in_n_interval in range(self.n_steps):
                if is_episode_done: # 이전 스텝에서 에피소드가 종료되었다면 루프 중단
                    break

                # 4.1 로컬 네트워크를 사용하여 행동 선택
                action_probabilities, state_value_prediction = self.local_network(current_state_tensor)
                
                action_distribution = torch.distributions.Categorical(action_probabilities)
                chosen_action = action_distribution.sample() # 행동 샘플링
                
                action_log_prob = action_distribution.log_prob(chosen_action) # 선택된 행동의 로그 확률
                policy_entropy = action_distribution.entropy() # 정책의 엔트로피

                # 4.2 선택된 행동을 환경에 적용
                next_observation_stacked, reward, terminated, truncated, _ = self.env.step(chosen_action.item())
                is_episode_done = terminated or truncated # 에피소드 종료 여부

                current_episode_reward += reward # 현재 에피소드 보상 누적

                # 4.3 경험(transition)을 로컬 버퍼에 저장
                log_probs_buffer.append(action_log_prob)
                values_buffer.append(state_value_prediction)
                rewards_buffer.append(torch.FloatTensor([reward])) # 보상을 (1,) 형태의 텐서로 저장
                entropies_buffer.append(policy_entropy)
                
                # 4.4 다음 상태 준비 및 글로벌 스텝 카운터 증가
                current_state_tensor = self.convert_obs_to_tensor(next_observation_stacked)
                
                with self.global_counter.get_lock(): # 공유 변수 접근 시 락 사용
                    self.global_counter.value += 1
                
                # 4.5 영상 녹화 중이면 현재 프레임 저장
                if trigger_video_recording:
                    rendered_rgb_frame = self.env.render()
                    if rendered_rgb_frame is not None:
                        frames_for_current_video.append(rendered_rgb_frame)
                
                if is_episode_done: # 에피소드 종료 시 n-스텝 루프도 중단
                    break
            
            # 5. n-스텝 후 또는 에피소드 종료 시 손실 계산 및 글로벌 네트워크 업데이트
            if log_probs_buffer: # 수집된 경험이 있을 경우에만 업데이트
                total_loss = self.calculate_n_step_loss(rewards_buffer, values_buffer, log_probs_buffer, entropies_buffer, is_episode_done)

                # 5.1 글로벌 옵티마이저의 그라디언트 초기화
                self.global_optimizer.zero_grad()
                # 5.2 로컬 네트워크의 그라디언트 계산 (loss.backward())
                total_loss.backward()
                
                # 5.3 (선택적) 로컬 네트워크의 그라디언트 클리핑
                torch.nn.utils.clip_grad_norm_(self.local_network.parameters(), 40.0) # 최대 norm 값은 조절 가능
                
                # 5.4 로컬 그라디언트를 글로벌 네트워크의 그라디언트로 복사
                for local_param, global_param in zip(self.local_network.parameters(),
                                                     self.global_network.parameters()):
                    if local_param.grad is not None: # 로컬 그라디언트가 존재하면
                        global_param._grad = local_param.grad # 글로벌 파라미터의 그라디언트로 할당
                                                              # (주의: _grad는 직접 접근, grad는 public attribute)

                # 5.5 글로벌 네트워크 파라미터 업데이트
                self.global_optimizer.step()

            # 6. 에피소드 종료 시 처리
            if is_episode_done:
                self.total_episodes_by_worker += 1
                print(f"워커 {self.worker_id}, 글로벌 스텝: {self.global_counter.value}, 에피소드 {self.total_episodes_by_worker} 종료, 보상: {current_episode_reward:.2f}")

                # 영상 저장 (녹화 중이었다면)
                if trigger_video_recording and frames_for_current_video:
                    timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
                    video_file_name = os.path.join(self.video_dir,
                                                  f"pong_w{self.worker_id}_ep{self.total_episodes_by_worker}_gs{self.global_counter.value}_{timestamp_str}.mp4")
                    save_video(frames_for_current_video, video_file_name)
                    frames_for_current_video = [] # 리스트 비우기
                    trigger_video_recording = False # 플래그 리셋

        self.env.close() # 워커 프로세스 종료 시 환경 정리
        print(f"워커 {self.worker_id} 종료.")

# %% [markdown]
# ## 7. 메인 학습 실행 함수
# 전체 학습 과정을 관리하는 메인 함수를 정의합니다. 글로벌 네트워크, 옵티마이저, 카운터를 생성하고 워커들을 실행시킵니다.

# %%
def run_a3c_training():
    """A3C 학습을 위한 메인 함수"""
    training_start_time = time.time()
    
    # 환경 정보 임시로 가져오기 (액션 수 등)
    # FrameStackObservation을 거치면 obs.shape[0]이 채널 수가 됨
    temp_env_for_info = create_env()
    num_input_channels = temp_env_for_info.observation_space.shape[0]
    num_actions = temp_env_for_info.action_space.n
    temp_env_for_info.close()
    print(f"환경 정보: 입력 채널 수 = {num_input_channels}, 액션 수 = {num_actions}")

    # 글로벌 공유 신경망 생성 및 공유 메모리 설정
    shared_global_network = ActorCriticNetwork(num_input_channels, num_actions)
    shared_global_network.share_memory() # 멀티프로세싱 환경에서 파라미터 공유

    # 글로벌 공유 옵티마이저 (Adam 사용)
    global_optimizer = optim.Adam(shared_global_network.parameters(), lr=LEARNING_RATE)

    # 글로벌 공유 스텝 카운터
    global_step_counter = mp.Value('i', 0) # 'i'는 정수형을 의미, 초기값 0

    # 워커 프로세스들 생성
    worker_processes = []
    for i in range(N_WORKERS):
        worker = A3CWorker(shared_global_network, global_optimizer, global_step_counter,
                             worker_process_id=i,
                             video_save_dir=VIDEOS_DIR,
                             record_video_interval=RECORD_INTERVAL_GLOBAL_STEPS,
                             max_global_training_steps=MAX_GLOBAL_STEPS)
        worker_processes.append(worker)

    print(f"{N_WORKERS}개의 워커 프로세스를 시작합니다...")
    for worker_proc in worker_processes:
        worker_proc.start() # 각 워커 프로세스 시작 (run_worker_process 메소드 실행)

    # --- 메인 프로세스에서 진행 상황 모니터링 및 주기적 모델 저장 ---
    last_model_saved_step = 0
    model_save_interval = 100000 # 예: 10만 글로벌 스텝마다 모델 저장 (값 조절 가능)
    
    monitoring_loop_active = True
    try:
        while monitoring_loop_active:
            time.sleep(30) # 30초마다 현재 상태 확인
            current_global_steps = global_step_counter.value
            elapsed_training_time = time.time() - training_start_time
            print(f"진행 상황: 글로벌 스텝 = {current_global_steps}/{MAX_GLOBAL_STEPS}, 경과 시간 = {elapsed_training_time:.2f}초")

            # 주기적 모델 저장 로직
            if current_global_steps >= last_model_saved_step + model_save_interval:
                model_save_path = os.path.join(MODELS_DIR, f"a3c_pong_checkpoint_gs{current_global_steps}.pt")
                torch.save(shared_global_network.state_dict(), model_save_path)
                print(f"체크포인트 모델이 저장되었습니다: {model_save_path}")
                last_model_saved_step = current_global_steps

            # 최대 글로벌 스텝 도달 시 모니터링 종료
            if current_global_steps >= MAX_GLOBAL_STEPS:
                print("최대 글로벌 스텝에 도달했습니다. 학습 및 모니터링을 중지합니다.")
                monitoring_loop_active = False
                break 
            
            # 모든 워커가 (예상치 못하게) 먼저 종료되었는지 확인
            all_workers_have_finished = True
            for w_p in worker_processes:
                if w_p.is_alive(): # 하나라도 살아있으면 아직 진행 중
                    all_workers_have_finished = False
                    break
            if all_workers_have_finished and current_global_steps < MAX_GLOBAL_STEPS:
                print("모든 워커가 최대 스텝 도달 전에 종료되었습니다. 모니터링을 중지합니다.")
                monitoring_loop_active = False # 루프 종료

    except KeyboardInterrupt: # Ctrl+C 입력 시
        print("사용자에 의해 학습이 중단되었습니다. 워커 프로세스를 종료합니다...")
    finally:
        # 모든 워커 프로세스가 종료될 때까지 대기 또는 강제 종료
        print("워커 프로세스 종료 처리 중...")
        for worker_proc_to_join in worker_processes:
            if worker_proc_to_join.is_alive():
                try:
                    worker_proc_to_join.terminate() # 강제 종료 시그널 전송
                    # terminate 후 join으로 좀비 프로세스 방지
                    worker_proc_to_join.join(timeout=10) # 10초 동안 종료 대기
                except Exception as e_join_terminate:
                    print(f"워커 {worker_proc_to_join.worker_id} 종료 중 예외 발생: {e_join_terminate}")
                if worker_proc_to_join.is_alive():
                    print(f"경고: 워커 {worker_proc_to_join.worker_id}가 정상적으로 종료되지 않았습니다.")
        print("모든 워커 프로세스에 대한 종료 시도가 완료되었습니다.")

    # 최종 모델 저장
    final_model_save_path = os.path.join(MODELS_DIR, f"a3c_pong_final_gs{global_step_counter.value}.pt")
    torch.save(shared_global_network.state_dict(), final_model_save_path)
    print(f"최종 모델이 저장되었습니다: {final_model_save_path}")

    total_training_duration = time.time() - training_start_time
    print(f"총 학습 시간: {total_training_duration:.2f}초")
    
    # 가상 디스플레이 중지 (Colab에서 필요)
    if 'display' in globals() and isinstance(display, Display) and display.is_started:
        try:
            display.stop()
            print("가상 디스플레이가 중지되었습니다.")
        except Exception as e_display_stop:
            print(f"가상 디스플레이 중지 중 오류: {e_display_stop}")

# %% [markdown]
# ## 8. 스크립트 실행
# `if __name__ == "__main__":` 블록을 사용하여 스크립트가 직접 실행될 때 `run_a3c_training()` 함수를 호출합니다.
# Colab에서는 이 셀을 직접 실행하면 됩니다. 멀티프로세싱 시작 방법 설정은 OS에 따라 중요할 수 있습니다.

# %%
if __name__ == "__main__":
    # 멀티프로세싱 시작 방법 설정 (특히 macOS나 Windows에서 'spawn' 또는 'forkserver'가 필요할 수 있음)
    # Colab (Linux 기반)에서는 'fork'가 기본값이지만, 'spawn'이 더 안정적일 수 있음
    # force=True는 이미 설정된 경우에도 강제로 재설정 시도
    try:
        # Colab에서는 mp.set_start_method를 최상위 수준에서 한 번만 호출하는 것이 좋음
        # (노트북 셀 실행 시 __main__ 블록이 매번 실행되므로 주의)
        # 이미 설정되었다면 RuntimeError 발생 가능
        if mp.get_start_method(allow_none=True) is None: # 아직 설정되지 않았다면
             mp.set_start_method('spawn', force=True)
             print("멀티프로세싱 시작 방법을 'spawn'으로 설정했습니다.")
        else:
            print(f"멀티프로세싱 시작 방법이 이미 '{mp.get_start_method()}'(으)로 설정되어 있습니다.")

    except RuntimeError as e_mp_start_method:
        print(f"멀티프로세싱 시작 방법 설정 중 오류: {e_mp_start_method} (이미 설정되었거나 지원되지 않을 수 있음)")
        # 이 오류가 발생해도 진행 가능할 수 있음
        pass
        
    run_a3c_training()

# %% [markdown]
# ## 9. 실행 가이드 및 추가 정보
#
# ### 주요 기능
# - 멀티프로세싱을 통한 비동기 학습 (A3C)
# - Actor-Critic 아키텍처 사용
# - CNN 기반 특징 추출 (Atari 환경용)
# - 학습 과정 영상 저장 (Worker 0 담당)
# - 주기적인 모델 체크포인트 저장
#
# ### 수동 설치 및 실행 (로컬 환경)
# 1.  **Python 환경 구성**: Python 3.8 이상 권장.
# 2.  **필요한 패키지 설치**:
#     ```bash
#     pip install torch gymnasium[atari,accept-rom-license] ale-py numpy opencv-python tqdm matplotlib imageio pyvirtualdisplay
#     ```
#     또는 제공된 `requirements.txt` 파일이 있다면:
#     ```bash
#     pip install -r requirements.txt
#     ```
# 3.  **스크립트 실행**:
#     ```bash
#     python your_script_name.py
#     ```
#
# ### 주요 하이퍼파라미터 (스크립트 상단에서 조절 가능)
# - `MAX_GLOBAL_STEPS`: 총 학습 스텝 수 (기본값: 200,000)
# - `RECORD_INTERVAL_GLOBAL_STEPS`: 영상 녹화 간격 (기본값: 50,000)
# - `N_WORKERS`: 병렬로 실행할 워커 수 (기본값: 4)
# - `LEARNING_RATE`: 옵티마이저 학습률 (기본값: 1e-4)
# - `GAMMA`: 할인 계수 (기본값: 0.99)
# - `ENTROPY_BETA`: 정책 엔트로피 정규화 가중치 (기본값: 0.01)
# - `N_STEPS_UPDATE`: n-스텝 업데이트를 위한 스텝 수 (기본값: 20)
#
# ### 출력 디렉토리 구조
# 스크립트 실행 시 `output_a3c_pong_colab` (또는 `OUTPUT_DIR`에 지정된 이름) 디렉토리가 생성되며 내부는 다음과 같습니다:
# ```
# output_a3c_pong_colab/
# ├── models/          # 학습된 모델 (.pt 파일) 저장
# ├── logs/            # (현재 코드에서는 사용 안함, 필요시 추가 가능)
# └── videos/          # 학습 과정 중 녹화된 영상 (.mp4 파일) 저장
# ```
#
# ### 주의사항
# 1.  Colab 환경에서는 CPU 코어 수와 RAM 제한을 고려하여 `N_WORKERS` 및 `MAX_GLOBAL_STEPS`를 조절하는 것이 좋습니다.
# 2.  학습 시간은 하이퍼파라미터 설정 및 Colab 세션에 할당된 하드웨어 성능에 따라 크게 달라질 수 있습니다.
# 3.  Colab에서 생성된 파일들은 세션이 종료되면 사라질 수 있으므로, 중요한 결과물(모델, 영상)은 Google Drive에 마운트하여 저장하거나 로컬로 다운로드하는 것이 좋습니다.