Bắt đầu với TorchRL cho Deep Reinforcement Learning

Thiết lập TorchRL

Trong phần này, tôi sẽ hướng dẫn bạn cách cài đặt và bắt đầu với TorchRL. 

Yêu cầu tiên quyết

Bạn cần một vài gói phần mềm phụ thuộc trước khi cài đặt và sử dụng TorchRL:

• PyTorch: TorchRL được xây dựng trên PyTorch, vì vậy bạn cần cài đặt trước. 
• Gymnasium: Bạn cần gói Gymnasium để import các môi trường RL. Tính đến tháng 1/2025, phiên bản mới nhất của Gymnasium không tương thích với TorchRL, như giải thích trên trang Git Discussions này. Do đó, hãy cài đặt phiên bản cũ hơn 0.29.1. 
• PyGame: Đây là một gói Python cho trò chơi điện tử.  TorchRL cần nó để mô phỏng các môi trường RL dạng trò chơi, như CartPole. 
• TensorDict: Việc dùng cấu trúc dữ liệu dạng từ điển để lưu trữ đầu vào và đầu ra của mạng nơ-ron giúp thao tác với tensor trong TorchRL thuận tiện hơn. TensorDict cung cấp một vùng chứa tensor lưu trữ các tensor theo cặp khóa-giá trị. 

Cài đặt các gói tiên quyết:

!pip install torch tensordict  gymnasium==0.29.1 pygame

Cài đặt TorchRL

Cài đặt TorchRL bằng pip. Tôi rất khuyến nghị cài đặt các gói này trong môi trường Conda nếu bạn làm việc trên máy cá nhân hoặc máy chủ. 

!pip install torchrl

Xác minh cài đặt

Sau khi cài đặt, hãy kiểm tra TorchRL đã được cài đặt thành công. Thử import torchrl trong một shell (hoặc notebook) Python. Dùng phương thức check_env_specs() để kiểm tra một môi trường tiêu chuẩn (như CartPole) có phù hợp với đặc tả của torchrl hay không:

import torchrl 
from torchrl.envs import GymEnv 
from torchrl.envs.utils import check_env_specs

check_env_specs(GymEnv("CartPole-v1"))

Kết quả đầu ra sẽ cho biết môi trường đã được tạo thành công và hoạt động với TorchRL. 

[torchrl][INFO] check_env_specs succeeded!

Các thành phần chính của TorchRL

Trước khi xây dựng tác tử RL đầu tiên của bạn, hãy xem các khối xây dựng cốt lõi của TorchRL.

Môi trường

TorchRL cung cấp một API thống nhất để giao tiếp với các môi trường khác nhau. Nó thực hiện điều này bằng cách bọc các chức năng đặc thù của từng môi trường vào một tập hợp tiêu chuẩn các lớp và hàm bao. Bạn truyền các tham số phù hợp cho lớp bao, và nội bộ nó sẽ ánh xạ lệnh của bạn tới lời gọi hàm tương ứng của môi trường cụ thể. Cụ thể:

• TorchRL chuyển đổi các trạng thái môi trường (quan sát), hành động và phần thưởng thành các đối tượng tensor của PyTorch, có thể được dùng trực tiếp bởi các mô-đun triển khai thuật toán RL.
• Nó cho phép áp dụng các bước tiền xử lý và hậu xử lý, ví dụ như chuẩn hóa và co giãn các tensor đầu vào hoặc trình bày các tensor đầu ra theo một định dạng cụ thể.

Ví dụ, để tạo một môi trường từ Gymnasium, hãy dùng mô-đun GymEnv:

env = GymEnv("CartPole-v1")

Biến đổi (Transforms)

Việc mở rộng môi trường gốc với các tính năng bổ sung và các phép biến đổi mà nếu không bạn sẽ phải tự triển khai là rất phổ biến. Ví dụ, bạn có thể thêm mô-đun đếm bước vào môi trường thay vì tự viết bộ đếm. Bộ đếm bước theo dõi số bước trong mỗi tập. Mô-đun transformedEnv hỗ trợ việc này:

from torchrl.envs import GymEnv, StepCounter, TransformedEnv
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())

Tương tự, chuẩn hóa tensor trước khi thao tác là một bước tiền xử lý phổ biến. Một phép biến đổi chuẩn hóa, sử dụng mô-đun ObservationNorm sẽ chuẩn hóa các tensor. Tài liệu của hàm TransformedEnv trình bày các loại biến đổi khác nhau. 

Bạn cũng có thể kết hợp nhiều biến đổi bằng tham số compose():

base_env = GymEnv('CartPole-v1', device=device) 

env = TransformedEnv( 
    base_env, 
    Compose(
        ObservationNorm(in_keys=["observation"]), 
        StepCounter()
    )
)

Tác tử và chính sách

Trong RL, tác tử quyết định hành động của mình dựa trên chính sách và trạng thái quan sát được của môi trường. Mục tiêu của tác tử là tối đa hóa tổng phần thưởng nhận được từ môi trường. Nó nhận phần thưởng khi chọn đúng hành động (chẳng hạn như dừng lại khi đèn đỏ). 

Ví dụ, trong một chiếc xe tự lái, tác tử là người điều khiển xe. Các quyết định của nó (như đánh lái theo một hướng nhất định) dựa trên quan sát về trạng thái môi trường (như giao thông, vị trí của xe, v.v.) và chính sách của nó (như tránh người đi bộ và các chướng ngại vật khác, dừng ở đèn đỏ, v.v.). 

Chính sách đơn giản nhất là chọn hành động một cách ngẫu nhiên. Diễn viên (actor) sẽ ngẫu nhiên chọn một trong các hành động từ không gian hành động có thể. Đôi khi chính sách ngẫu nhiên được dùng để tạo bộ dữ liệu tương tác ban đầu trước khi bắt đầu huấn luyện mô hình. 

Dùng mô-đun RandomPolicy để tạo một chính sách ngẫu nhiên. Chính sách ngẫu nhiên này nhận tham số action_spec chỉ định không gian hành động. Trong ví dụ dưới đây, không gian hành động gồm các số (liên tục) giữa -1 và +1. Vì vậy, chính sách ngẫu nhiên sẽ chọn một hành động được biểu diễn bởi một số ngẫu nhiên trong khoảng -1 đến +1. 

import torchrl 
import torch
from tensordict import TensorDict
from torchrl.data.tensor_specs import Bounded

action_spec = Bounded(-torch.ones(1), torch.ones(1))
actor = torchrl.envs.utils.RandomPolicy(action_spec=action_spec) 
td = actor(TensorDict({}, batch_size=[])) 
print(td.get("action"))

Đầu ra sẽ là một tensor, như dưới đây:

tensor([0.9258])

Kiểm tra chính sách thực sự là ngẫu nhiên bằng cách chạy actor một lần nữa:

td = actor(TensorDict({}, batch_size=[])) 
print(td.get("action"))

Nó sẽ in ra một tensor ngẫu nhiên khác.

Nếu bạn cần tài liệu nhập môn về các khái niệm RL, hãy xem lộ trình kỹ năng Reinforcement Learning in Python trên DataCamp!

Xây dựng tác tử RL đầu tiên với TorchRL

Trong phần này, tôi hướng dẫn bạn cách triển khai một tác tử Reinforcement Learning đơn giản bằng TorchRL. 

Trước khi bắt đầu, hãy import các gói phần mềm tiên quyết trong Python:

• time để đo thời gian huấn luyện tác tử.
• GymEnv, StepCounter và TransformedEnv để làm việc với các môi trường Gymnasium.
• MLP để tạo một mạng nơ-ron MLP (perceptron nhiều lớp) đơn giản. 
• EGreedyModule để cân bằng giữa khám phá môi trường và khai thác chính sách tốt nhất đã biết.
• QValueModule và DQNLoss để triển khai thuật toán Deep Q-Learning.
• SoftUpdate để cập nhật mạng nơ-ron.
• SyncDataCollector để thu thập dữ liệu từ các tương tác của tác tử. 
• ReplayBuffer để lưu trữ dữ liệu từ các tương tác của tác tử.
• Adam cho lan truyền ngược.
• matplotlib để trực quan hóa tiến trình huấn luyện. 
• torchrl_logger để ghi log phiên huấn luyện.

import time
import matplotlib.pyplot as plt
from torchrl.envs import GymEnv, StepCounter, TransformedEnv
from tensordict.nn import TensorDictModule as TensorDict, TensorDictSequential as Seq
from torchrl.modules import EGreedyModule, MLP, QValueModule
from torchrl.objectives import DQNLoss, SoftUpdate
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer
from torch.optim import Adam
from torchrl._utils import logger as torchrl_logger

Bước 1: Định nghĩa môi trường

Trong ví dụ này, chúng ta giải quyết môi trường CartPole. Import môi trường này từ Gymnasium cùng với bộ đếm bước để theo dõi số bước huấn luyện:

env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())

Gieo hạt (seed) cho Python và các môi trường RL để tái lập kết quả tương tự giữa các phiên huấn luyện. 

torch.manual_seed(0)
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())
env.set_seed(0)

Định nghĩa các tham số và siêu tham số cho việc huấn luyện: 

• INIT_RAND_STEPS: Số bước mà tác tử hành động ngẫu nhiên trước khi dùng chính sách. Những bước đầu này dùng để thu thập dữ liệu ban đầu nhằm khởi động việc huấn luyện chính sách.
• FRAMES_PER_BATCH: Số điểm dữ liệu (mỗi điểm tương ứng một tương tác hoặc một bước thời gian) trong một batch huấn luyện. 
• OPTIM_STEPS: Số bước tích lũy loss trước khi chạy một lượt lan truyền ngược. 
• EPS_0: Giá trị epsilon ban đầu, hệ số khám phá. 
• BUFFER_LEN: Kích thước replay buffer.
• ALPHA: Tốc độ học.
• TARGET_UPDATE_EPS: Hệ số suy giảm để cập nhật mạng đích bằng mô-đun soft-update.
• REPLAY_BUFFER_SAMPLE: Kích thước mẫu ngẫu nhiên lấy từ replay buffer ở mỗi vòng lặp huấn luyện. Replay buffer lưu kết quả tương tác của tác tử với môi trường ở mỗi bước thời gian. 
• LOG_EVERY: Số bước sau đó sẽ in tiến độ huấn luyện.
• MLP_SIZE: Kích thước của mạng nơ-ron.

INIT_RAND_STEPS = 5000 
FRAMES_PER_BATCH = 100
OPTIM_STEPS = 10
EPS_0 = 0.5
BUFFER_LEN = 100_000
ALPHA = 0.05
TARGET_UPDATE_EPS = 0.95
REPLAY_BUFFER_SAMPLE = 128
LOG_EVERY = 1000
MLP_SIZE = 64

Bước 2: Tạo chính sách

Định nghĩa một mạng nơ-ron đơn giản để triển khai chính sách: 

• Định nghĩa MLP (mạng nơ-ron). Với một quan sát, MLP xuất ra giá trị hành động, thực hiện vai trò của hàm Q. 
• Định nghĩa một từ điển tensor bằng mô-đun TensorDict trên MLP ở trên. Từ điển này ánh xạ các quan sát trạng thái của môi trường (khóa) tới các giá trị xác suất hành động (giá trị) tương ứng với trạng thái đó. 
• Dùng QValueModule để triển khai hàm Q. Cho một tensor chứa các giá trị hành động, nó trả về hành động có giá trị cao nhất. Nó triển khai chiến lược tham lam (khai thác). 
• Kết hợp (dùng TensorDictSequential module) từ điển tensor của MLP và QValueModule để định nghĩa chính sách. 

value_mlp = MLP(out_features=env.action_spec.shape[-1], num_cells=[MLP_SIZE, MLP_SIZE])
value_net = TensorDict(value_mlp, in_keys=["observation"], out_keys=["action_value"])
policy = Seq(value_net, QValueModule(spec=env.action_spec))

• Định nghĩa hàm khám phá bằng EGreedyModule - nó nhận đầu vào là không gian hành động của môi trường, độ dài replay buffer và hệ số khám phá epsilon. Xâu chuỗi (dùng mô-đun TensorDictSequential) hàm khám phá này với chính sách đã định nghĩa ở trên để có chính sách cuối cùng:

exploration_module = EGreedyModule(
    env.action_spec, annealing_num_steps=BUFFER_LEN, eps_init=EPS_0
)
policy_explore = Seq(policy, exploration_module)

Bước 3: Huấn luyện tác tử

Bước đầu tiên trong huấn luyện tác tử là thu thập dữ liệu từ các tương tác của tác tử với môi trường. Dùng SyncDataCollector để xây dựng bộ thu thập chạy chính sách và thu thập kết quả tương tác của tác tử: 

collector = SyncDataCollector(
    env,
    policy_explore,
    frames_per_batch=FRAMES_PER_BATCH,
    total_frames=-1,
    init_random_frames=INIT_RAND_STEPS,
)

Tạo một replay buffer để lưu kết quả các tương tác:

rb = ReplayBuffer(storage=LazyTensorStorage(BUFFER_LEN))

Bạn cũng cần khai báo các mô-đun dành riêng cho huấn luyện, như hàm mất mát (để dùng cách tính loss dựa trên DQN), bộ tối ưu (thuật toán Adam truyền thống) và các hàm cập nhật (để cập nhật mạng nơ-ron). Tất cả đều dựa trên các mô-đun dựng sẵn của TorchRL:

loss = DQNLoss(value_network=policy, action_space=env.action_spec, delay_value=True)
optim = Adam(loss.parameters(), lr=ALPHA)
updater = SoftUpdate(loss, eps=TARGET_UPDATE_EPS)

Khởi tạo các bộ đếm để theo dõi tổng số bước và số tập, số bước thành công mỗi tập, và thời gian thực thi:

total_count = 0
total_episodes = 0
t0 = time.time()
success_steps = []

Hàm huấn luyện gồm 2 vòng lặp for: 

• Vòng ngoài cùng chạy chính sách và thêm kết quả tương tác của tác tử vào replay buffer.
• Vòng trong chia các mẫu huấn luyện thành các batch. Nó xử lý mỗi batch như sau: 
• Chọn một mẫu ngẫu nhiên từ replay buffer.
• Tính loss bằng hàm mất mát.
• Chạy bộ tối ưu và hàm cập nhật.
• Cập nhật các bộ đếm.

Mỗi bước đều gọi các mô-đun dựng sẵn của TorchRL mà không cần viết mã từ đầu. Đoạn mã dưới đây cho thấy cách triển khai các bước này:

for i, data in enumerate(collector):
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > INIT_RAND_STEPS:
        for _ in range(OPTIM_STEPS):
            sample = rb.sample(REPLAY_BUFFER_SAMPLE)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()

            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())

            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    success_steps.append(max_length)

Bước 4: Đánh giá tác tử

Vòng lặp ở phần trước liên tục huấn luyện chính sách. Chúng ta cần đặt tiêu chí đánh giá hiệu năng và quyết định khi nào coi việc huấn luyện là thành công. Chúng ta cũng muốn in ra tiến độ huấn luyện. 

Chúng ta in tiến độ huấn luyện theo chu kỳ bằng logger của TorchRL:

if total_count > 0 and total_count % LOG_EVERY == 0:
    torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")

Chúng ta dùng số bước tối đa mà tác tử đạt được trong tập gần nhất để xác định liệu huấn luyện đã thành công hay chưa. Môi trường CartPole-v1 giới hạn số bước tối đa và tổng phần thưởng mỗi tập ở 500. Thông thường, coi một chính sách là thành công nếu đạt hơn 475 bước:

if max_length > 475:
    print("TRAINING COMPLETE")
    break

Hai đoạn mã trên nên được thêm vào cuối vòng lặp huấn luyện (đã trình bày trước đó). Đoạn dưới đây cho thấy vòng lặp huấn luyện kèm mã đánh giá tác tử: 

for i, data in enumerate(collector):
    # Write data in replay buffer
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > INIT_RAND_STEPS:
        for _ in range(OPTIM_STEPS):
            sample = rb.sample(REPLAY_BUFFER_SAMPLE)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()

            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())

            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    success_steps.append(max_length)

    if total_count > 0 and total_count % LOG_EVERY == 0:
        torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")

    if max_length > 475:
        print("TRAINING COMPLETE")
        break

Cuối cùng, sau khi huấn luyện và đánh giá, in tổng thời gian huấn luyện và vẽ biểu đồ tiến trình huấn luyện:

t1 = time.time()

torchrl_logger.info(
    f"solved after {total_count} steps, {total_episodes} episodes and in {t1-t0}s."
)

def plot_steps():
    plt.plot(success_steps)
    plt.title('Successful steps over training episodes')
    plt.xlabel('Training episodes')
    plt.ylabel('Steps')
    plt.show()

plot_steps()

Bạn có thể tìm và chạy toàn bộ mã để triển khai DQN bằng TorchRL trong sổ tay DataLab này. 

Trong phần này, chúng ta đã xây dựng và huấn luyện một tác tử RL bằng thuật toán DQN đơn giản. Ở phần tiếp theo, chúng ta khám phá cách dùng các mô-đun dựng sẵn của TorchRL để triển khai một thuật toán phức tạp hơn như PPO. 

Khám phá các thuật toán dựng sẵn trong TorchRL

Như đã đề cập ở các phần trước, TorchRL đi kèm một vài thuật toán dựng sẵn. Hãy xem chúng và cách chúng hoạt động.

Các thuật toán được hỗ trợ

TorchRL bao gồm các mô-đun dựng sẵn cho nhiều thuật toán Học tăng cường sâu phổ biến, như: 

• Deep Q networks (DQN)
• Deep deterministic policy gradient (DDPG)
• Soft actor-critic (SAC)
• Randomized ensembled double Q-learning (REDQ)
• CrossQ
• Implicit Q-learning (IQL)
• Continuous Q-learning (CQL)
• Generative adversarial imitation learning (GAIL)
• Decision transformer (DT)
• Twin-delayed DDPG (TD3) 
• Actor-advantage critic (A2C)
• Proximal policy optimization (PPO)
• REINFORCE
• và nhiều hơn nữa 

Điều này giúp bạn dễ dàng thử nghiệm các loại thuật toán khác nhau và nghiên cứu hiệu năng của từng thuật toán trong việc giải một bài toán cụ thể. 

Ví dụ: PPO trong TorchRL

Proximal Policy Optimization sử dụng một hàm mục tiêu xấp xỉ có cắt (clipped surrogate objective) để đạt được quá trình huấn luyện mượt mà, cân bằng giữa khám phá và khai thác. 

Trong phần này, tôi sẽ trình bày cách dùng mô-đun PPO dựng sẵn trong TorchRL. Tương tự như triển khai trước đó, hãy import các mô-đun tiên quyết. Bên cạnh các mô-đun (như ReplayBuffer, SyncDataCollector, v.v.) cần cho triển khai trước, bạn cần thêm một vài gói cho PPO:

• ProbabilisticActor, để chọn hành động một cách ngẫu nhiên có kiểm soát. Việc chọn hành động ngẫu nhiên (thay vì luôn theo chính sách nhằm tối đa hóa phần thưởng) giúp tác tử khám phá môi trường và tìm ra những lộ trình tối ưu hơn trong quá trình huấn luyện.
• OneHotCategorical, để tạo mã hóa one-hot cho tensor biểu diễn log xác suất của các hành động. 
• ValueOperator, để xây dựng một mô-đun TorchRL dựa trên mạng nơ-ron triển khai hàm giá trị.
• GAE, để triển khai ước lượng lợi thế tổng quát (generalized advantage estimate). PPO dùng một hàm lợi thế như một đại diện cho hàm giá trị. Mô-đun giá trị (ở trên) là đầu vào cho GAE. 
• ClipPPOLoss, để xây dựng hàm mục tiêu có cắt khi triển khai PPO. 

import torch
from torch import nn

from torchrl.envs import Compose, ObservationNorm, DoubleToFloat, StepCounter, TransformedEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.utils import check_env_specs
from torchrl.modules import ProbabilisticActor, OneHotCategorical, ValueOperator
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.storages import LazyTensorStorage
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.objectives.value import GAE
from torchrl.objectives import ClipPPOLoss
from tensordict.nn import TensorDictModule

Khai báo các tham số và siêu tham số cho việc huấn luyện:

• FRAMES_PER_BATCH: Số khung hình (bước thời gian) trong mỗi batch dữ liệu từ tương tác của tác tử với môi trường. 
• SUB_BATCH_SIZE: Số bước thời gian trong mỗi sub-batch. Trong huấn luyện PPO, mỗi batch được chia thành các sub-batch. 
• TOTAL_FRAMES: Tổng số bước thời gian để chạy huấn luyện. 
• GAMMA: Hệ số chiết khấu cho phần thưởng ở các bước tương lai.
• CLIP_EPSILON: Xác định vùng tin cậy trong đó chính sách có thể thay đổi ở mỗi vòng lặp huấn luyện tiếp theo. Nó được thực hiện bằng cách cắt (clip) chính sách cập nhật sao cho tỷ lệ giữa xác suất mới và cũ nằm trong một giới hạn nhất định. 
• ALPHA: Tốc độ học.
• ENTROPY_EPS: Siêu tham số này điều khiển tỷ lệ giữa khám phá và khai thác. 
• OPTIM_STEPS: Số lần chạy bộ tối ưu trên mỗi sub-batch dữ liệu. 
• LOG_EVERY: Số bước sau đó sẽ đánh giá chính sách và in phần thưởng. 

FRAMES_PER_BATCH = 1024
TOTAL_FRAMES = 1048576
GAMMA = 0.99
LAMBDA = 0.95
CLIP_EPSILON = 0.2
ALPHA = 1e-4
ENTROPY_EPS = 5e-4 
SUB_BATCH_SIZE = 64
OPTIM_STEPS = 8
LOG_EVERY = 16

Import môi trường CartPole gốc từ Gymnasium:

device="cpu"
base_env = GymEnv('CartPole-v1', device=device) 

Khai báo môi trường TorchRL bằng cách import môi trường gốc và thêm các mô-đun để chuẩn hóa các tensor quan sát và đếm số bước:

env = TransformedEnv( 
    base_env, 
    Compose(
        ObservationNorm(in_keys=["observation"]), 
        DoubleToFloat(), 
        StepCounter()
    )
)

Khởi tạo môi trường và gieo hạt cho các bộ sinh số ngẫu nhiên:

env.transform[0].init_stats(1024) 
torch.manual_seed(0)
env.set_seed(0)
check_env_specs(env) 

Khai báo actor là một mạng nơ-ron với 1 lớp ẩn có 16 trọng số. Nó nhận đầu vào là quan sát từ môi trường và xuất ra xác suất của mỗi hành động trong không gian hành động. 

actor_net = nn.Sequential(
    nn.Linear(env.observation_spec["observation"].shape[-1], 32, device=device),
    nn.ReLU(),
    nn.Linear(32, 32, device=device),
    nn.ReLU(),
    nn.Linear(32, env.action_spec.shape[-1], device=device),
    nn.ReLU()
)

Tạo một từ điển tensor (dựa trên mạng nơ-ron ở trên) với các khóa và giá trị ánh xạ trạng thái tới xác suất hành động. 

actor_module = TensorDictModule(actor_net, in_keys=["observation"], out_keys=["logits"])

Đối với việc huấn luyện PPO (và nhiều thuật toán RL khác), bạn không phải lúc nào cũng chọn hành động dẫn tới phần thưởng cao nhất ở bước tiếp theo. Việc có một mức độ ngẫu nhiên khi chọn hành động cho phép tác tử khám phá môi trường và tìm ra những lộ trình tốt hơn có thể dẫn tới lợi ích dài hạn cao hơn. Tạo một actor xác suất để làm điều này:

actor = ProbabilisticActor(
    module = actor_module,
    spec = env.action_spec,
    in_keys = ["logits"],
    distribution_class = OneHotCategorical, 
    return_log_prob = True
)

Tạo mạng để triển khai hàm giá trị. Mạng này có 16 trọng số. Nó nhận đầu vào là trạng thái (quan sát) của môi trường và xuất ra giá trị kỳ vọng của trạng thái đó. 

value_net = nn.Sequential(
    nn.Linear(env.observation_spec["observation"].shape[-1], 16, device=device),
    nn.ReLU(),
    nn.Linear(16, 1, device=device),
    nn.ReLU()
)

Tạo một mô-đun TorchRL bao bọc hàm giá trị (đã triển khai ở trên) bằng ValueOperator(). Mô-đun này giao tiếp với các thành phần TorchRL khác. 

value_module = ValueOperator(
    module = value_net,
    in_keys = ["observation"]
)

Tạo replay buffer để lưu kết quả tương tác của tác tử với môi trường: 

replay_buffer = ReplayBuffer(
    storage = LazyTensorStorage(max_size=FRAMES_PER_BATCH),
    sampler = SamplerWithoutReplacement()
)

Tạo bộ thu thập dữ liệu để chạy chính sách trên môi trường và thu thập kết quả tương tác của tác tử với môi trường ở mỗi bước thời gian (frame):

collector = SyncDataCollector(
    env,
    actor,
    frames_per_batch = FRAMES_PER_BATCH,
    total_frames = TOTAL_FRAMES,
    split_trajs = True,
    reset_at_each_iter = True,
    device=device
)

Dùng mô-đun GAE để triển khai hàm lợi thế cho PPO. Hàm lợi thế dựa trên hàm giá trị. 

advantage_module = GAE(
    gamma = GAMMA, 
    lmbda = GAMMA, 
    value_network = value_module,
    average_gae = True
) 

Dùng mô-đun dựng sẵn ClipPPOLoss để triển khai hàm mất mát theo thuật toán PPO:

loss_module = ClipPPOLoss(
    actor_network = actor,
    critic_network = value_module,
    clip_epsilon = CLIP_EPSILON,
    entropy_bonus = bool(ENTROPY_EPS),
    entropy_coef = ENTROPY_EPS
)

Khai báo bộ tối ưu Adam: 

optim = torch.optim.Adam(loss_module.parameters(), lr=ALPHA)

Tạo một scheduler để giảm dần tốc độ học khi quá trình huấn luyện diễn ra: 

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, TOTAL_FRAMES // FRAMES_PER_BATCH)

Xây dựng vòng lặp huấn luyện bằng các thành phần đã tạo ở trên. Hàm huấn luyện gồm ba vòng lặp for:

• Vòng ngoài cùng chạy qua dữ liệu được thu thập từ các tương tác của tác tử với môi trường. Tham số TOTAL_FRAMES quyết định chạy tác tử bao nhiêu bước thời gian. 
• Vòng giữa thực hiện các việc sau: 
• Chạy vòng trong cùng với số lần cố định (thường từ 5 đến 10), được xác định bởi siêu tham số OPTIM_STEPS. 
• Cập nhật tốc độ học sau mỗi vòng huấn luyện. 
• Tính giá trị hàm lợi thế cho mỗi bước thời gian. 
• Định kỳ đánh giá chính sách và in phần thưởng. 
• Vòng trong cùng chạy vòng lặp huấn luyện:
• Lấy mẫu một batch dữ liệu huấn luyện từ replay buffer.
• Tính loss bằng mô-đun mất mát. 
• Chạy lan truyền ngược trên loss.
• Dùng bộ tối ưu để thực hiện gradient descent. 

Đoạn mã dưới đây triển khai vòng lặp huấn luyện:

for i, tensordict_data in enumerate(collector): 
    for _ in range(OPTIM_STEPS): 
        advantage_module(tensordict_data)
        replay_buffer.extend(tensordict_data.reshape(-1).cpu())
        for _ in range(FRAMES_PER_BATCH // SUB_BATCH_SIZE): 
            data = replay_buffer.sample(SUB_BATCH_SIZE)
            loss = loss_module(data.to(device))
            loss_value = loss["loss_objective"] + loss["loss_critic"] + loss["loss_entropy"]
            loss_value.backward()
            optim.step()
            optim.zero_grad()
    scheduler.step()

    if i % LOG_EVERY == 0:
        with torch.no_grad():
            rollout = env.rollout(FRAMES_PER_BATCH, actor)
            reward_eval = rollout["next","reward"].sum()
            #print(rollout["next","reward"].sum())
            print(reward_eval)
            rewards.append(reward_eval)
            del rollout

Sổ tay DataLab này có mã (như trên) để huấn luyện một tác tử RL bằng thuật toán PPO từ TorchRL. Hãy dùng nó làm điểm khởi đầu để tinh chỉnh tham số và cải thiện hiệu năng của tác tử. 

Tùy biến thuật toán

TorchRL có thiết kế dạng mô-đun cho phép tùy biến các giải pháp RL. Khung làm việc chịu trách nhiệm kết nối giữa các thành phần. Để kiểm thử và so sánh hiệu năng, bạn có thể hoán đổi môi trường, các chính sách, replay buffer và thuật toán RL. Nhờ tính mô-đun này:

• Bạn có thể thay thế các thành phần khác nhau mà không cần viết lại phần lớn chương trình. 
• Từng thành phần cũng có thể được chỉnh sửa độc lập để tạo ra các giải pháp tùy chỉnh. 

• Bạn có thể thử nghiệm với các mô-đun khác nhau và xây dựng một giải pháp tối ưu. 

Một số ví dụ về tùy biến trong TorchRL gồm:

• Vòng lặp huấn luyện có thể được tùy biến với scheduler tốc độ học, logger và các chỉ số tùy chỉnh. 
• Có thể dùng các độ dài replay buffer khác nhau để điều chỉnh thời lượng huấn luyện.
• Môi trường có thể được tùy biến để bao gồm các hàm co giãn và chuẩn hóa, bộ đếm bước, v.v. 

Trực quan hóa và gỡ lỗi quá trình huấn luyện RL

Các phần trước cho thấy cách dùng các mô-đun TorchRL để huấn luyện tác tử RL. Trong phần này, chúng ta thảo luận cách theo dõi và trực quan hóa tiến trình huấn luyện. 

Theo dõi tiến trình huấn luyện

Việc ghi log nhiều chỉ số trong quá trình huấn luyện sẽ hữu ích để theo dõi tiến độ. Các gói như TensorBoard giúp bạn dễ dàng ghi log kết quả trực tiếp trong quá trình huấn luyện. Ví dụ, SummaryWriter của TensorBoard cho phép bạn gọi trong vòng lặp huấn luyện để ghi log các chỉ số. Mã giả sau minh họa cách làm:

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir="training_logs")

for _ in training_data (): # training loop
    loss = … 
    loss.backprop() 
    writer.add_scalar(“Loss - “, loss.item(), step)
   … 

writer.close() 

Sau khi hoàn tất huấn luyện và ghi log, bạn có thể vẽ và trực quan hóa kết quả. Ví dụ, dùng lệnh tensorboard để trực quan hóa tiến trình huấn luyện. 

tensorboard --logdir=”training_logs”

Gỡ lỗi với TorchRL

Gỡ lỗi là cần thiết để xác minh tương tác của tác tử với môi trường. Bước đầu tiên, bạn muốn kiểm tra không gian hành động và không gian quan sát. Tài liệu của môi trường nên bao gồm các đặc tả cho môi trường. Ví dụ, tài liệu của CartPole cho thấy: 

• Không gian hành động gồm hai giá trị rời rạc - 0 (đẩy xe sang trái) và 1 (đẩy xe sang phải) 
• Không gian quan sát có 4 giá trị: 
• Vị trí của xe đẩy, với giá trị trong khoảng 
• -4,8 đến +4,8. 
• Vận tốc của xe đẩy, với bất kỳ giá trị số thực nào
• Góc của cột, với giá trị trong khoảng -24 ° đến  +24 °
• Vận tốc góc của cột, với bất kỳ giá trị số thực nào

Hãy kiểm tra không gian quan sát và hành động để xem chúng có tương ứng với các giá trị kỳ vọng hay không. Ví dụ:

print("Observation space: ", env.observation_spec)
print("Action space: ", env.action_spec)

Ngoài ra, bạn có thể lấy các mẫu ngẫu nhiên từ không gian quan sát và hành động và kiểm tra giá trị của chúng: 

print("Sample Observation:", base_env.observation_spec.sample().get("observation"))
print("Sample Action:", base_env.action_spec.sample())

Đầu ra tương tự mẫu dưới đây:

Sample Observation: tensor([-4.5960e+00,  3.4028e+38,  2.2261e-02,  3.4028e+38])
Sample Action: tensor([1, 0])

Lưu ý rằng 4 giá trị trạng thái nằm trong các khoảng đã nêu trước đó. Lặp lại các lệnh trên vài lần và để ý rằng các giá trị đầu ra (lấy mẫu ngẫu nhiên) khác nhau đều nằm trong các khoảng chỉ định. 

Hai lệnh trên dựa trên môi trường gốc, được import trực tiếp từ Gymnasium. Trên thực tế, chúng ta sử dụng môi trường đã biến đổi, nơi áp dụng các biến đổi như chuẩn hóa lên giá trị tensor gốc. Do đó, khi dùng môi trường TorchRL đã biến đổi để lấy mẫu quan sát, các giá trị của chúng có thể không còn nằm trong cùng khoảng như ở môi trường Gymnasium gốc. 

Ví dụ, lấy mẫu từ môi trường đã biến đổi:  

print("Sample Observation:", env.observation_spec.sample().get("observation"))

Đầu ra tương tự: 

Sample Observation: tensor([-57.5869,      nan,   4.5276,      nan])

Lưu ý rằng các giá trị vị trí xe đẩy và góc của cột nằm ngoài các khoảng đã nêu. Điều này là do các giá trị tensor đã được chuẩn hóa. Hơn nữa, vận tốc của xe đẩy và vận tốc góc của cột có giá trị nan vì tác tử chưa bắt đầu tương tác với môi trường. 

Trực quan hóa hiệu năng của tác tử

Ngoài việc vẽ đồ thị tiến trình huấn luyện, việc render môi trường và quan sát trực quan các tương tác của tác tử cũng hữu ích. Điều này có thể mang lại hiểu biết về cách tác tử đang hoạt động. 

Cách thực dụng nhất để trực quan hóa môi trường là dựng video. Bạn cần một vài gói bổ sung: 

• torchvision, để làm việc với tệp đa phương tiện
• av, để bao ffmpeg nhằm có thể giao tiếp với PyTorch. 

pip install torchvision
pip install av==12.0.0

Sau khi hoàn tất huấn luyện, hãy thực hiện các bước chuẩn bị sau để dựng video:

• Khai báo đường dẫn (tương đối) để lưu video.
• Khởi tạo logger để lưu kết quả tương tác của tác tử ở định dạng CSV.
• Khởi tạo bộ ghi video để dùng các log CSV tạo video.
• Biến đổi môi trường TorchRL để thêm bộ ghi video.

…
# training loop
# for _ in …
…
…

from torchrl.record import CSVLogger, VideoRecorder 

path = "./training_loop"
logger = CSVLogger(exp_name="dqn", log_dir=path, video_format="mp4")
video_recorder = VideoRecorder(logger, tag="video")
record_env = TransformedEnv(
    GymEnv("CartPole-v1", from_pixels=True, pixels_only=False), video_recorder
)

Chạy chính sách (đã huấn luyện) trên môi trường và ghi các khung hình render vào tệp video:

record_env.rollout(max_steps=1000, policy=policy)
video_recorder.dump()

Sau khi chạy chương trình, bạn sẽ thấy video trong thư mục (./training_loop trong đoạn trên) mà bạn đã đặt trước đó. Lưu ý rằng sổ tay DataLab triển khai DQN bằng TorchRL không bao gồm mã tạo video vì việc xuất tệp video với notebook trực tuyến là khó. Hãy chạy chương trình (và thêm các bước trên) trên máy tính cục bộ hoặc máy chủ để ghi video.

Thực hành tốt nhất khi sử dụng TorchRL

Cuối cùng, hãy điểm qua một số thực tiễn tốt nhất. Dưới đây là các khuyến nghị của tôi.

Bắt đầu với các môi trường đơn giản

Mặc dù TorchRL giúp việc phát triển trở nên thuận tiện hơn, huấn luyện các tác tử RL đạt hiệu suất tốt trong môi trường phức tạp vẫn là một thách thức. Do đó, trước khi thử giải các bài toán khó, điều thiết yếu là hãy giải quyết các môi trường đơn giản như CartPole. 

Tính mô-đun của TorchRL cho phép bạn thử nghiệm với các thuật toán và tham số khác nhau. Trước khi áp dụng các thuật toán thay thế và tùy biến trong thực tế, cần có một số hiểu biết về cách chúng vận hành. Điều này hiệu quả nhất trong môi trường đơn giản, nơi tác động của từng thay đổi riêng lẻ dễ quan sát một cách trực quan. 

Thử nghiệm với siêu tham số

Hiệu suất huấn luyện của các mô hình RL nhạy với nhiều siêu tham số, như tốc độ học, tỷ lệ thăm dò, hệ số chiết khấu, và các siêu tham số đặc thù thuật toán khác, như tỷ lệ cắt cho PPO. Ví dụ, tốc độ học quá cao sẽ khiến huấn luyện không ổn định, còn quá thấp sẽ khiến hội tụ quá lâu. 

Tương tự, tỷ lệ thăm dò quá cao sẽ ngăn huấn luyện hội tụ, còn quá thấp có thể khiến tác tử không khám phá được quỹ đạo tối ưu. 

Không có công thức cố định để xác định giá trị siêu tham số đúng. Có các hướng dẫn và giá trị khuyến nghị cho những môi trường chuẩn, nhưng chúng có thể không phù hợp cho mọi bài toán. Vì vậy, cần thử nghiệm bằng nhiều phương pháp, như tìm kiếm lưới hoặc tìm kiếm ngẫu nhiên, để xác định các giá trị siêu tham số tốt nhất. 

Bạn cũng có thể dùng các thư viện tự động như Weights & Biases Sweeps hoặc Optuna để kiểm tra hiệu suất huấn luyện trên nhiều tổ hợp siêu tham số khác nhau. 

Tận dụng các thuật toán dựng sẵn

Như bạn đã thấy trong các ví dụ trước, dùng các mô-đun dựng sẵn của TorchRL giúp tiết kiệm đáng kể công sức phát triển. Phương án thay thế là xây dựng mọi chức năng từ đầu, sẽ tiêu tốn nhiều thời gian và chi phí để phát triển và kiểm thử hơn rất nhiều. 

Vì nhiều nhà phát triển sử dụng các mô-đun của TorchRL nên chúng cũng đã được kiểm thử trong nhiều kịch bản khác nhau. Do đó, bạn có thể kỳ vọng chúng ít lỗi hơn so với mô-đun tự xây. Vì vậy, đối với hầu hết các trường hợp sử dụng tiêu chuẩn, rất nên dùng các mô-đun dựng sẵn. 

Kết luận

Trong bài viết này, chúng ta đã đề cập các khái niệm cơ bản của TorchRL, một khung công tác dựa trên PyTorch để triển khai các thuật toán RL. Chúng ta cũng đã xem các ví dụ thực hành sử dụng TorchRL để triển khai những lời giải đơn giản như Deep Q-Learning và các thuật toán phức tạp hơn như PPO. 

Bước tiếp theo, bạn có thể dùng các chương trình này làm nền tảng để thử nghiệm với những môi trường và thuật toán khác.

Nếu bạn muốn đào sâu hiểu biết về các nền tảng RL và ứng dụng thực tiễn, hãy xem khóa học Reinforcement Learning with Gymnasium in Python để có thêm trải nghiệm thực hành với các môi trường dựa trên Gymnasium.