GPU에서 강화학습을 돌리는 법: 환경과 학습을 텐서로 연결하기

강화학습 코드를 GPU에서 실행하려면 무엇을 바꿔야 할까? 가장 먼저 떠오르는 것은 model.to("cuda")다. 정책 신경망의 계산을 GPU로 옮기는 출발점이지만, 강화학습 프로그램 전체의 실행 속도는 신경망만으로 결정되지 않는다.

정책은 환경에서 관측을 받고, 행동을 고르고, 그 행동의 결과로 다음 관측과 보상을 얻는다. 이 경험을 모아야 비로소 학습할 수 있다. 작은 신경망을 쓰는 문제에서는 경험을 만드는 환경 실행이 학습보다 더 많은 시간을 차지할 수 있다. GPU가 행동을 계산하는 동안보다 CPU의 환경 처리를 기다리는 시간이 길다면, 모델만 GPU에 올려도 기대한 효과가 나지 않는다.

이 글에서는 환경 상태를 배치 텐서로 표현하고, 경험 수집과 정책 업데이트까지 연결하는 방법을 살펴본다. 설명에는 상태와 행동이 작은 턴제 게임을 사용한다. 중요한 것은 게임의 공략이 아니라, 다른 RL 문제에도 적용할 수 있는 데이터의 모양, 장치 간 이동, 종료 처리, 학습 표본의 정의다.

1. GPU를 어디에 사용할지 먼저 정한다

강화학습에서 GPU를 쓰는 구성은 크게 두 가지로 나눌 수 있다.

구성	환경 실행	정책 추론과 학습	주요 확인 사항
CPU 환경 + GPU 정책	CPU 프로세스 또는 벡터 환경	GPU	환경 처리량과 관측·행동 전송 비용
GPU 환경 + GPU 정책	GPU 텐서 연산	GPU	전이의 벡터화, 배치 크기와 메모리

이미 CPU용 게임 엔진이나 시뮬레이터가 있다면 첫 번째 구성이 자연스럽다. 여러 환경을 병렬로 실행하고, 관측을 모아 한 번에 GPU 정책에 전달한다. 이때도 CPU 환경 객체에 .to("cuda")를 붙인다고 내부 로직이 GPU 코드로 바뀌지는 않는다.

환경의 규칙을 텐서 연산으로 표현할 수 있다면 두 번째 구성을 만들 수 있다. 환경 상태, 관측, 행동, 학습 표본이 GPU에 머물도록 구성한다. 아래 예시는 이 경로를 따른다. 외부 게임 엔진, 파일 입출력, 복잡한 동적 자료구조에 의존하는 환경까지 이 방식으로 쉽게 옮길 수 있다는 뜻은 아니다.

GPU 안에서 N개 환경의 상태가 관측, 정책, 행동을 거쳐 다시 상태로 연결되고, 수집 버퍼가 PPO 업데이트에 쓰이는 흐름 — 정책은 수집 구간 동안 고정하고, 모은 표본으로 업데이트한 뒤 다음 수집을 시작한다. CPU에는 배치 단위 지표와 체크포인트를 필요한 시점에 전달한다.

2. 환경 객체 N개를 상태 텐서로 바꾼다

CPU 구현에서는 환경마다 객체를 하나씩 만들기 쉽다. 각 객체가 현재 위치, 남은 자원, 종료 여부를 갖고, step() 메서드가 조건문으로 상태를 바꾼다.

GPU 구현에서는 같은 종류의 상태를 하나의 텐서에 모은다. 각 행이 독립된 환경 하나를 나타내도록 만드는 것이다. 예를 들어 환경이 4,096개라면 현재 위치를 저장하는 텐서의 모양은 [4096]이고, 각 환경의 두 플레이어가 가진 자원은 [4096, 2]로 표현한다.

예시 게임에는 여섯 칸 중 한 곳에 숨겨진 탄환이 있다. FIRE는 현재 칸을 사용하고, 빈칸이면 다음 칸과 상대 차례로 넘어간다. 플레이어마다 한 번씩 쓸 수 있는 PEEK은 현재 칸을 확인하고 차례를 유지한다. SKIP은 칸을 바꾸지 않고 상대에게 차례를 넘긴다. 탄환을 발사한 플레이어가 패배한다.

데이터	모양	의미
`chamber_position`	`[N]`	각 환경의 현재 칸
`current_player`	`[N]`	현재 행동할 플레이어
`peek_available`, `skip_available`	`[N, 2]`	플레이어별 남은 능력
`done`	`[N]`	종료된 환경을 표시하는 boolean mask
`observations`	`[N, 6]`	정책에 공개하는 관측
`actions`	`[N]`	환경별로 선택한 행동

환경의 전체 상태와 정책의 관측은 다르다. 환경은 판정을 위해 숨겨진 탄환 위치를 알아야 하지만, 정책에 그 값을 전달하면 문제가 바뀐다. 이 예시의 관측은 현재 칸, 자신과 상대의 남은 능력, PEEK으로 확인한 정보만 포함한다.

이 구분은 다른 문제에서도 같다. 시뮬레이터 내부의 정답이나 미래 정보를 정책 입력에 섞지 않아야 한다. 또한 여섯 관측값만 쓰는 feed-forward 정책이 모든 문제에 충분한 것은 아니다. 과거 관측을 기억해야 하는 문제라면 관측 이력이나 순환 신경망 등 별도의 설계가 필요하다.

3. 환경별 조건문을 마스크 연산으로 바꾼다

벡터화의 핵심은 for env in environments를 빠르게 돌리는 데 있지 않다. 어떤 행에 어떤 전이를 적용할지 boolean mask로 계산하는 데 있다.

아래는 첨부 코드의 step()에서 FIRE 처리에 관련된 부분을 추린 것이다. 전체 상태 갱신을 설명하는 발췌이며, 단독 실행용 코드는 아니다.

active = ~self.done
fire_mask = active & (actions == FIRE)
is_live = self.chamber_position == self.live_position

fire_live = fire_mask & is_live
fire_blank = fire_mask & ~is_live

self.chamber_position = (
    self.chamber_position + fire_blank.to(torch.long)
)
self.done = self.done | fire_live

fire_blank가 참인 행만 현재 칸을 한 칸 전진시킨다. 이미 끝난 환경은 active가 거짓이므로 다시 진행되지 않는다. 전체 구현은 같은 방식으로 플레이어 전환, 자원 소모, 관측 정보와 보상을 갱신한다.

이때 서로 연관된 마스크는 갱신 전 상태를 기준으로 먼저 계산해야 한다. 현재 플레이어를 먼저 바꾼 뒤 보상을 계산하면 실제로 행동한 사람이 아닌 상대에게 보상을 줄 수 있다. CPU 구현에서는 자연스럽게 보였던 순서가 벡터화 과정에서 쉽게 뒤섞인다.

torch.where(mask, next_value, old_value)도 자주 사용한다. 다만 Python의 if처럼 선택하지 않은 쪽의 계산을 생략하는 구문은 아니다. 전달하는 값들은 먼저 계산되므로, 선택되지 않을 것이라는 이유로 0으로 나누는 식 등을 넣어서는 안 된다.

시간 방향의 루프까지 없애야 하는 것은 아니다. 다음 상태는 현재 행동의 결과에 의존하므로 for step in range(T)는 남을 수 있다. 먼저 제거하려는 것은 같은 시점의 N개 환경을 하나씩 처리하는 Python 루프다.

4. 정책도 N개 관측을 한 번에 처리한다

환경이 [N, observation_dim] 관측을 만들면 정책은 이를 받아 [N, action_dim] logits를 출력한다. 별도의 정책 객체 N개가 필요한 것이 아니라, 하나의 정책이 배치 전체를 처리한다.

예시에서는 이미 사용한 PEEK이나 SKIP을 다시 고르지 않도록 행동을 마스킹한다. 다음 코드는 rl_gpu/policy.py의 계산 구조를 단순화한 것이다.

logits, values = policy(observations)

legal = torch.stack(
    (
        torch.ones_like(observations[:, 0], dtype=torch.bool),
        observations[:, 1] > 0,  # 자신의 PEEK이 남아 있는가?
        observations[:, 2] > 0,  # 자신의 SKIP이 남아 있는가?
    ),
    dim=1,
)

masked_logits = logits.masked_fill(~legal, -1e9)
log_probs = torch.log_softmax(masked_logits, dim=1)
actions = torch.multinomial(
    log_probs.exp(), 1, generator=policy_rng
).squeeze(1)

모든 행에 적어도 하나의 유효 행동이 있어야 한다. 여기서는 FIRE가 항상 가능하다. 연속 행동 문제라면 이런 범주형 마스크 대신 행동 범위와 확률분포를 다른 방식으로 정의해야 한다.

환경 상태, 관측, 정책 파라미터, 행동, 난수 생성기는 서로 맞는 장치에 있어야 한다. CUDA 관측을 CPU 모델에 넣거나 CPU generator를 CUDA 샘플링에 쓰는 식의 혼합을 피한다.

5. rollout을 고정된 모양으로 수집한다

정책이 환경에서 행동하며 모은 경험을 보통 rollout이라고 부른다. 일정 길이 T 동안 N개 환경을 실행하면 다음과 같은 버퍼를 만들 수 있다.

버퍼	모양	업데이트에 쓰는 정보
관측	`[T, N, observation_dim]`	행동을 고른 당시의 입력
행동	`[T, N]`	실제 선택한 행동
이전 log probability	`[T, N]`	수집 당시 정책의 선택 확률
이전 value	`[T, N]`	수집 당시의 가치 추정
return 또는 advantage	`[T, N]`	얼마나 좋은 결과였는가
학습 mask	`[T, N]`	손실에 포함할 의사결정인가

이 예시에서는 한 게임이 최대 10번의 행동 안에 끝난다. 따라서 모든 환경을 시작하고 10 step을 수집한 뒤, 완결 게임의 결과로 학습할 수 있다. 먼저 끝난 행은 그대로 유지하고 나머지 시간 칸을 padding으로 취급한다. 다음 rollout을 시작할 때 리셋한다.

이것은 짧고 종료 길이가 제한된 문제에서 선택한 수집 방식이다. 일반적인 긴 에피소드에서는 일정한 T만큼 수집한 뒤 계속 진행할 수 있다. 그 경우 실제 종결인 terminated와 시간 제한 등의 truncated를 구분하고, return 계산에 필요한 bootstrap을 처리해야 한다. 자동 리셋을 사용하는 환경에서는 마지막 상태의 관측을 새 에피소드의 첫 관측과 혼동하지 않도록 해야 한다.

수집 버퍼의 메모리는 대략 T × N × 저장할 특징 수에 비례한다. 작은 6차원 관측에서는 부담이 작지만, 이미지 관측이나 긴 시퀀스에서는 관측 버퍼가 모델보다 클 수도 있다. N을 무조건 크게 하기 전에 이 곱부터 계산하는 편이 좋다.

6. 무엇을 학습 표본으로 삼을지 명확히 한다

GPU에서 한꺼번에 계산할 수 있다는 사실과 모든 행을 학습에 써도 된다는 사실은 다르다.

턴제 예시에서는 learner와 상대가 번갈아 행동한다. learner의 정책을 학습할 때 상대의 행동을 learner가 선택한 행동처럼 취급하면 안 된다. 끝난 게임의 padding도 학습에서 제외해야 한다.

learner_turn = env.current_player == learner_seats
active = ~env.done
weights[step] = (active & learner_turn).float()

승패 역시 learner의 관점으로 계산한다. 마지막에 상대가 탄환을 발사했다면 learner에게는 승리다. 마지막 행동자의 음수 보상을 그대로 learner의 보상으로 복사하면 부호가 뒤집힌다.

손실 평균은 전체 [T, N] 원소 수가 아니라 유효한 learner 결정 수로 나눈다. 이 설계에서는 policy loss뿐 아니라 value loss, entropy, advantage 정규화에도 같은 mask를 사용한다. 의미 없는 padding이 평균이나 분산에 들어가면 정책이 같은 경험을 했어도 학습 결과가 달라질 수 있다.

단일 에이전트 문제에서는 상대 차례를 구분할 필요가 없지만, padding·종료·유효 표본을 구분해야 한다는 원칙은 그대로 남는다.

7. 수집과 PPO 업데이트를 연결한다

PPO는 수집 당시 정책과 현재 정책의 행동 확률 비율을 이용한다. 이 글의 구현은 clipped PPO objective를 사용한다.

r_t(\theta)=\exp\left(\log\pi_\theta(a_t\mid o_t)-\log\pi_{\mathrm{old}}(a_t\mid o_t)\right)

확률 비율과 advantage로 만든 목적함수를 clipping하고, value loss와 entropy 항을 함께 사용한다. 목적과 알고리즘의 원형은 PPO 논문에서 확인할 수 있다.

실행 순서는 다음과 같다.

정책을 고정한 채 torch.no_grad()에서 rollout을 모은다.
수집 당시의 log probability와 value를 보관한다.
return과 advantage를 계산한다. 예시에서는 완결 게임의 learner 승패를 사용한다.
[T, N, ...]를 [T × N, ...]로 펼치고 유효 표본 mask를 유지한다.
현재 정책으로 확률을 다시 계산하고 역전파한다.
정해진 횟수의 optimizer 업데이트를 마친 뒤 새 rollout을 수집한다.

수집은 gradient 없이, 업데이트는 gradient를 켠 상태에서 실행한다. 이전 log probability를 현재 정책의 값으로 매번 덮어쓰거나, 전체 학습 루프를 no_grad()로 감싸는 실수를 피해야 한다.

첨부 구현은 작은 관측과 네트워크를 전제로 전체 배치를 네 번 재사용한다. 더 큰 모델에서는 minibatch 분할이 필요할 수 있다. PPO라고 해서 오래된 rollout을 무제한 재사용하는 것은 아니다. 수집 정책과 업데이트된 정책이 벌어지는 정도를 함께 확인해야 한다. PyTorch의 PPO 튜토리얼에는 수집과 손실 계산을 구성하는 다른 예시도 있다.

8. 최소 실행부터 확인한다

예제 코드 묶음 다운로드에는 환경, 정책, rollout 수집, PPO 업데이트, CLI와 테스트가 들어 있다. 압축을 풀고 gpu-rl-example/ 안에서 실행한다. 코드를 읽는 순서는 envs.py → policy.py → rollout.py → optimization.py → experiment.py가 자연스럽다.

Python 3.12 환경에서 의존성을 설치한다. 아래 버전은 예제 검증에 사용한 조합이다. GPU 실행에는 PyTorch CUDA 런타임과 호환되는 NVIDIA 드라이버가 필요하다.

python3.12 -m venv .venv
.venv/bin/python -m pip install -r requirements.txt
.venv/bin/python -m unittest discover -v

다음은 수집 한 번과 PPO 업데이트 한 번을 연결한 실행 가능한 코드다. 압축 파일의 quickstart.py와 같다.

import torch

from rl_gpu.envs import RussianRouletteVectorized
from rl_gpu.policy import ActorCritic
from rl_gpu.rollout import collect_episodes
from rl_gpu.optimization import update_policy


def main() -> None:
    """Run one complete collection/update cycle on CUDA or CPU."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(11)
    env = RussianRouletteVectorized(4096, device=device, seed=12)
    policy = ActorCritic().to(env.device)
    optimizer = torch.optim.Adam(policy.parameters(), lr=0.001)
    policy_rng = torch.Generator(device=env.device).manual_seed(13)
    opponent_rng = torch.Generator(device=env.device).manual_seed(14)

    batch = collect_episodes(policy, env, policy_rng, opponent_rng)
    if batch.observations.device != env.device:
        raise RuntimeError("rollout changed device")
    if not bool(batch.weights.sum() > 0):
        raise RuntimeError("no learner decisions to optimize")
    metrics = update_policy(policy, optimizer, batch)
    print("device:", env.device)
    print("learner decisions:", int(batch.weights.sum()))
    print("loss:", metrics.loss)


if __name__ == "__main__":
    main()

.venv/bin/python quickstart.py

이 실행은 데이터 수집, 역전파, 파라미터 업데이트가 연결되는지 확인하는 단계다. 한 번 실행됐다는 사실만으로 정책이 잘 학습됐다고 판단하지 않는다.

반복 학습과 검증을 하려면 새 출력 경로를 지정한다.

.venv/bin/python -m rl_gpu.train \
  --output artifacts/my-first-run \
  --seeds 11 --num-envs 4096 \
  --min-updates 150 --eval-every 25 --max-updates 1000

이 명령은 환경 4,096개를 한 업데이트마다 완결 게임까지 진행한다. 150은 이 예제의 시작 설정이지 일반적인 수렴 기준이 아니다. CLI는 초기 정책 대비 승률 개선과 연속 검증 통과를 확인하고, 최고 체크포인트를 별도 평가한다. 최대 업데이트에 도달해도 목표를 못 채우면 실패로 종료한다. --device cpu를 지정하면 CPU에서도 같은 흐름을 확인할 수 있다.

저장한 모델을 다시 평가하는 단계까지 확인한다.

.venv/bin/python -m rl_gpu.evaluate \
  --checkpoint artifacts/my-first-run/seed-11/best.pt \
  --device cpu --episodes 24576

9. GPU가 바빠 보이는 것과 빨라진 것은 다르다

작은 N에서는 GPU 커널을 실행하고 계산을 준비하는 비용이 실제 연산보다 클 수 있다. 반대로 N을 늘리면 한 번의 연산으로 처리하는 일이 많아지지만, 메모리 사용량도 증가한다. 환경 처리량과 전체 학습 시간을 따로 측정해야 적절한 배치 크기를 고를 수 있다.

측정할 구간도 분리한다. 환경 step()만 측정하는지, 행동 샘플링과 리셋까지 포함하는지, 역전파까지 포함한 한 업데이트를 측정하는지에 따라 숫자의 의미가 달라진다. 같은 조건을 여러 번 반복하고 N도 함께 기록한다.

CUDA 연산은 기본적으로 비동기 실행된다. CPU 시계만 읽으면 GPU 작업이 끝나기 전에 타이머가 멈출 수 있다. 실제 경과 시간을 재려면 측정 구간 앞뒤에서 동기화하거나 CUDA event를 사용해야 한다. 이 동작은 PyTorch CUDA 문서에 설명되어 있다.

반대로 매 환경 step마다 .item(), .cpu(), .numpy()로 결과를 꺼내거나 CUDA 텐서를 Python 조건문으로 검사하면 CPU가 결과를 기다리는 지점이 생길 수 있다. 지표 출력은 배치가 끝난 뒤처럼 필요한 경계에 모은다. 다만 수치 오류 검출 등 필요한 검사를 없애면서까지 동기화를 피할 필요는 없다.

고정 길이 rollout에는 종료 후 padding이 포함될 수 있다. 따라서 T × N은 버퍼 칸 수이며 항상 실제 환경 전이 수와 같지는 않다. 처리량을 비교할 때는 어느 쪽을 세는지 명시해야 한다.

10. 빠른 코드가 올바른 학습을 한다는 보장은 없다

벡터화한 환경은 먼저 작은 CPU 기준 구현과 비교한다. 같은 seed를 넣는 것만으로 CPU와 CUDA의 난수열이 같아진다고 가정하지 말고, 동일한 초기 상태와 동일한 행동을 직접 주입해 전이를 대조한다.

적어도 다음 항목은 분리해 확인할 필요가 있다.

환경: 종료된 행이 다시 움직이지 않는가? 리셋할 행만 초기화되는가? 보상이 중복 지급되지 않는가?
관측: 정책에 숨겨진 정답이 들어가지 않는가? 실제 배포 시에도 얻을 수 있는 정보인가?
학습: 상대 행동과 padding이 손실에 섞이지 않는가? optimizer가 실제 파라미터를 바꾸는가? loss와 gradient가 유한한가?
평가: 초기 정책과 같은 조건으로 비교하는가? 체크포인트 선택에 쓴 검증과 최종 평가를 구분했는가? 한 좌석이나 특정 상대에만 유리하지 않은가?

GPU는 같은 시간에 더 많은 경험을 처리할 수 있게 해준다. 잘못된 보상이나 관측을 고쳐 주지는 않으며, 더 많은 에피소드가 계속 더 좋은 정책을 만든다는 보장도 없다. 학습 곡선이 정체되면 실행 시간을 늘리기 전에 환경 규칙, 상대 구성, 탐색, 정책 입력을 살펴봐야 한다.

처음부터 최대 배치를 목표로 삼을 필요는 없다. 작은 CPU 환경에서 규칙을 확인하고, N개 상태를 텐서로 묶고, rollout과 손실의 모양을 맞춘 뒤, 실제 병목에 맞춰 GPU 사용 범위를 넓힌다. 이 순서를 지키면 GPU 최적화가 학습 문제를 가리는 대신, 검증된 학습 과정을 더 빠르게 반복하는 도구가 된다.