[강화학습 정복하기] 2강: 강화학습을 위한 PyTorch 기초: 텐서부터 오토그라드까지

 

이 글은 서강대 가상융합 전문대학원(구 메타버스 전문대학원) 강화학습 자료를 기반으로 작성되었습니다.

 

2강: 강화학습을 위한 PyTorch 기초: 텐서부터 오토그라드까지

강화학습은 환경 속에서 에이전트가 보상을 극대화하도록 학습하는 과정이다. 이 학습이 가능하려면, 에이전트가 세상을 수치적으로 이해하고 경험을 데이터로 처리할 수 있는 기반이 필요하다. 그 중심에 있는 기술이 바로 PyTorch이다. PyTorch는 딥러닝을 구현하기 위한 핵심 도구로, 데이터를 표현하는 Tensor, 기울기를 자동으로 계산하는 Autograd, 신경망 구조를 설계하는 nn.Module, 그리고 학습 과정을 체계적으로 관리하는 Training Pipeline을 제공한다. 이 네 가지 구성 요소는 딥러닝 모델의 동작을 가능하게 하는 기반이며, 이러한 구조를 이해하는 것은 강화학습에서 정책 네트워크나 가치 네트워크를 설계하고 학습시키는 과정을 이해하는 데 도움이 된다.

 

딥러닝에서 모든 정보는 Tensor로 표현된다. Tensor는 단순한 숫자의 집합이 아니라, 데이터를 계산 가능한 형태로 구조화한 수학적 객체이다. 스칼라(0차원 값), 벡터(값의 나열), 행렬(2차원 데이터), 그리고 이미지나 시퀀스 같은 고차원 구조까지 모두 Tensor로 다룰 수 있다. torch.tensor(), zeros(), ones(), rand() 등의 함수로 쉽게 생성할 수 있으며, GPU를 활용하면 대규모 연산도 빠르게 수행된다. 결국 Tensor는 에이전트가 상태(state), 행동(action), 보상(reward)을 수치로 인식하도록 만드는 기반 언어이다. Tensor의 진정한 힘은 연산의 확장성에 있다. PyTorch의 브로드캐스팅(broadcasting)은 서로 다른 크기의 Tensor 간 연산을 자동으로 확장해 처리한다. reshape(), permute(), squeeze(), unsqueeze() 같은 기능은 데이터 구조를 자유롭게 변환할 수 있게 해준다. 이 덕분에 수많은 상태와 행동 데이터를 벡터화(batch) 연산으로 동시에 처리할 수 있다. 이러한 기술은 강화학습에서 여러 환경을 병렬로 시뮬레이션하거나, 정책 네트워크가 다중 상태를 한 번에 평가할 때 매우 유용하다. 학습의 안정성을 위해서는 Tensor를 효율적으로 다루는 습관도 중요하다. setup_seed(42)로 난수를 고정하면 실험의 재현성이 확보되고, torch.no_grad()는 추론 시 불필요한 미분 계산을 막아 메모리 낭비를 줄인다. 또한 .contiguous()를 통해 메모리 상에서 Tensor의 배열을 정렬하면 처리 속도가 개선된다. 이런 작은 습관들이 학습의 일관성과 효율성을 높이는 핵심이다. Tensor는 단순한 데이터 구조를 넘어 수학적 사고를 코드로 표현하는 도구다. mean(), std(), quantile()을 사용해 보상의 평균이나 분산을 계산할 수 있고, torch.bmm()으로 여러 배치(batch)의 행렬 연산을 병렬 처리할 수 있다. 또한 eig()와 svd() 같은 선형대수 함수는 모델의 수치적 안정성을 높이고, 특징 공간을 단순화하는 데 활용된다.

 

강화학습에서는 이러한 Tensor 연산이 실제 학습 수식으로 이어진다. 예를 들어,
returns = rewards + γ * next_values * (1 - done_mask) 는 벨만 기대 방정식(Bellman Expectation Equation)을 코드로 구현한 형태이다. 이 식을 통해 에이전트는 과거의 보상과 미래의 기대값을 통합하여 행동의 가치를 평가한다. 결국 Tensor는 강화학습의 이론적 기반을 실제 동작으로 바꾸는 핵심 매개체다.

딥러닝이 가능해진 이유 중 하나는 자동미분(Autograd) 기능 덕분이다. Autograd는 사용자가 정의한 연산 과정을 추적해 계산 그래프를 만들고, loss.backward() 한 줄로 모든 기울기를 자동 계산한다. 이 기울기(gradient)는 모델이 손실을 최소화하기 위해 어떤 방향으로 이동해야 하는지를 알려주는 신호다. 즉, 사람이 일일이 미분 공식을 구현하지 않아도 PyTorch가 자동으로 학습의 핵심 절차를 수행한다. Autograd는 강화학습의 정책 기울기(policy gradient) 개념과 직결된다. 정책 네트워크는 보상을 최대화하는 방향으로 파라미터를 조정해야 하는데, 이때 Autograd는 보상 신호로부터 기울기를 추적해 정책을 개선하는 역할을 한다. 또한 PyTorch는 기울기의 흐름을 정교하게 제어할 수 있는 기능을 제공한다. detach()는 특정 부분의 계산 그래프 연결을 끊어 불필요한 역전파를 막고, clip_grad_norm_()은 기울기가 지나치게 커지는 폭주를 방지한다. torch.no_grad()는 추론 단계에서 미분 기록을 완전히 차단해 효율을 높인다.

이러한 제어 기법들은 강화학습의 Actor–Critic 구조처럼 두 네트워크가 상호작용할 때 안정성을 유지하는 데 특히 중요하다. 고차 미분(create_graph=True) 기능은 또 다른 차원의 응용을 가능하게 한다. 이는 메타러닝(meta-learning), 즉 ‘학습하는 법을 학습하는 모델’을 구현할 때 사용된다. 강화학습에서는 에이전트가 단순히 보상을 극대화하는 수준을 넘어, 학습 전략 자체를 최적화하는 방식으로 발전할 수 있다.

 

nn.Module은 신경망의 설계도를 구성하는 핵심 도구다. 사용자는 이 클래스를 상속받아 입력층, 은닉층, 출력층을 정의하고, forward() 함수를 통해 데이터가 네트워크를 따라 흘러가는 과정을 설계한다. 이 방식은 수식을 직접 다루는 대신, 데이터의 흐름을 코드로 구현하는 접근법이다. 강화학습의 Q-network나 Policy-network 역시 이 구조를 기반으로 설계된다. 모델의 표현력은 층(layer)과 활성화 함수(activation function)의 선택에 따라 달라진다. Linear, Conv2d, LSTM, BatchNorm, Dropout 등을 적절히 조합해 다양한 데이터 형태를 처리할 수 있다. 활성화 함수로는 ReLU, Tanh, GELU, SiLU 등이 사용되며, 이들은 비선형성을 추가하여 모델이 복잡한 패턴을 학습하도록 돕는다. 모델이 학습하는 방향은 손실 함수(loss)와 최적화 알고리즘(optimizer)에 의해 결정된다. MSELoss, CrossEntropyLoss는 예측 오차를 계산하고, Adam, SGD, RMSprop 같은 옵티마이저가 기울기를 따라 파라미터를 갱신한다. 학습률 스케줄러(StepLR, CosineAnnealingLR)는 학습 강도를 시간에 따라 조절해 안정적인 수렴을 유도한다. 이는 강화학습에서 정책 안정화나 탐색–활용 균형을 유지하는 메커니즘과 유사하다. 가중치 초기화(Xavier, He)와 정규화(BatchNorm, Dropout, L2)는 학습의 균형을 유지하는 핵심 요소다. 초기화는 모델의 수렴 속도에 영향을 주고, 정규화는 과적합을 방지하여 일반화 성능을 높인다. 강화학습에서는 이러한 설정이 탐험 안정성(exploration stability)을 보장하는 역할을 한다.

 

학습 루프는 일련의 리듬을 가진 과정으로 구성된다. optimizer.zero_grad()로 초기화 후, 순전파 → 손실 계산 → 역전파 → 최적화 단계를 반복하며 모델이 점차 개선된다. evaluate() 단계에서는 torch.no_grad()를 사용해 모델의 성능을 점검한다. 이 흐름은 강화학습의 에피소드 단위 학습과 본질적으로 동일한 구조다. 마지막으로, 학습된 모델은 torch.save()로 저장하고 필요할 때 load_state_dict()로 불러 사용할 수 있다. 혼합정밀도 학습(GradScaler)과 torch.compile()은 GPU 자원을 절약하고 실행 속도를 높이는 실용적 방법이다. 이러한 기술은 실제 강화학습 실험에서 긴 학습 시간을 줄이고 안정성을 확보하는 데 큰 도움이 된다.

 

정리하면, Tensor로 데이터를 수치적으로 표현하고, Autograd로 학습 원리를 구현하며, nn.Module로 신경망을 설계하는 이 딥러닝 구조는, 강화학습의 두뇌를 만들기 위한 강력한 '구현 도구(Tool)'를 제공한다. 강화학습 고유의 이론(벨만 방정식, 정책 기울기 등 '설계도')을 이해하고 있는 사람이 이 '도구'를 활용할 때, 정책과 가치 함수를 코드로 구현하고 보상 신호를 기반으로 스스로 학습하는 에이전트를 완성할 수 있다.