Deep Q-Learning od podstaw: jak zbudować agenta RL z RLax, JAX i Haiku

Uczenie przez wzmacnianie (reinforcement learning) to jeden z tych obszarów AI, który wciąż potrafi zaskoczyć. Nie dlatego, że jest nowy, ale dlatego, że ekosystem narzędzi wokół niego zmienia się szybciej, niż większość z nas nadąża.

Opublikowany właśnie tutorial pokazuje, jak zbudować kompletnego agenta Deep Q-Learning (DQN) dosłownie od zera, korzystając z biblioteki RLax stworzonej przez Google DeepMind oraz zestawu narzędzi opartych na JAX. Brzmi technicznie? Jest. Ale warto zrozumieć, co kryje się pod maską.

Czym właściwie jest RLax i po co nam to?

RLax to biblioteka badawcza Google DeepMind zaprojektowana specjalnie do budowania algorytmów uczenia przez wzmacnianie z użyciem JAX. Kluczowe słowo to „badawcza” – RLax nie podaje gotowych, spakowanych algorytmów. Zamiast tego dostarcza zestaw wielokrotnie używalnych elementów, z których programista sam składa własny pipeline treningowy.

W tym tutorialu RLax współpracuje z trzema innymi bibliotekami:

• JAX – do wydajnych obliczeń numerycznych i automatycznego różniczkowania
• Haiku – do definicji i zarządzania sieciami neuronowymi
• Optax – do optymalizacji gradientowej

Razem tworzą coś w rodzaju zestawu klocków, z których można złożyć właściwie dowolny algorytm RL. To podejście zupełnie inne niż gotowe frameworki pokroju Stable Baselines.

CartPole jako plac testowy

Środowiskiem treningowym jest CartPole z biblioteki Gymnasium – klasyczny benchmark w świecie RL. Zadanie jest proste: utrzymać kij w równowadze na wózku, przesuwając go w lewo lub prawo. Prosto brzmi, ale agent musi sam odkryć tę strategię, ucząc się wyłącznie z sygnału nagrody.

To właśnie urok reinforcement learning. Nie mówimy agentowi „zrób tak” – on sam musi to wywnioskować.

Jak wygląda architektura agenta?

Sieć neuronowa zdefiniowana przez Haiku to prosta sieć MLP z dwiema ukrytymi warstwami po 128 neuronów każda. Wyjście to wektor wartości Q dla każdej możliwej akcji. W CartPole mamy tylko dwie: lewo i prawo.

Cały pipeline treningowy opiera się na kilku kluczowych komponentach:

• Replay buffer – przechowuje do 50 000 przejść (obserwacja, akcja, nagroda, następna obserwacja)
• Epsilon-greedy exploration – agent zaczyna od eksploracji losowej, stopniowo przechodząc do eksploatacji wyuczonej wiedzy
• Target network – oddzielna sieć aktualizowana miękko (soft update z parametrem tau), która stabilizuje uczenie
• Huber loss – funkcja straty odporna na duże błędy TD (temporal difference)

Piotr Wolniewicz, Redaktor Naczelny AIPORT.pl:

Takie tutoriale budzą we mnie mieszane uczucia. Z jednej strony – świetnie, że środowisko naukowe pokazuje, jak budować agenty RL od podstaw, bez czarnych skrzynek. Rozumienie mechanizmów, a nie tylko wywoływanie gotowych funkcji, to fundament, bez którego trudno o prawdziwy postęp badawczy. Z drugiej strony muszę zapytać: dla kogo to jest? Programista, który dopiero zaczyna przygodę z RL, może się pogubić w gąszczu JAX, Haiku, Optax i RLax jednocześnie. A doświadczony badacz pewnie już to wie. Brakuje mi w tym ekosystemie warstwy pośredniej – narzędzi, które są zarówno przejrzyste, jak i dostępne. Wartość edukacyjna tego tutorialu jest wysoka, ale droga do jego pełnego zrozumienia jest stroma. Czy to dobrze? Zależy, kogo pytasz.

Jak przebiega trening?

Trening trwa 40 000 kroków. Przez pierwsze 1000 kroków agent tylko zbiera doświadczenia do bufora (faza rozgrzewki). Potem co cztery kroki wykonuje krok treningowy: pobiera losową próbkę z bufora, oblicza błędy TD z użyciem prymitywu q_learning z biblioteki RLax, wylicza stratę Hubera, oblicza gradienty i aktualizuje wagi sieci.

Co 2000 kroków agent jest oceniany na osobnym środowisku ewaluacyjnym – bez eksploracji losowej, czysta eksploatacja. Wyniki są zapisywane i wizualizowane po zakończeniu treningu.

Modularność jako filozofia

Najciekawszym aspektem tego podejścia jest filozofia modularności. RLax nie narzuca struktury algorytmu. Dostarcza prymitywów – takich jak q_learning czy huber_loss – które można łączyć w dowolny sposób. Dzięki temu przejście z DQN do bardziej zaawansowanych metod, jak Double DQN, dystrybucyjny RL czy metody actor-critic, wymaga zmiany zaledwie kilku linii kodu.

To z kolei otwiera drogę do eksperymentowania. Zmień architekturę sieci. Podmień optymalizator. Przetestuj inne funkcje nagrody. Cały pipeline pozostaje czytelny i modyfikowalny.

Gdzie znaleźć kod?

Pełny notebook z kodem jest dostępny na GitHubie w repozytorium Marktechpost: github.com/Marktechpost/AI-Tutorial-Codes-Included

Wystarczy zainstalować zależności jedną komendą:

pip install "jax[cpu]" dm-haiku optax rlax gymnasium matplotlib numpy

I można zacząć. Środowisko CartPole nie wymaga GPU, więc całość można uruchomić nawet na zwykłym laptopie.

RL wraca do łask?

Warto zadać sobie pytanie, dlaczego akurat teraz pojawia się tyle materiałów o reinforcement learning. Odpowiedź jest dosyć prosta: modele językowe oparte na RLHF (Reinforcement Learning from Human Feedback) i ich kolejne ewolucje, jak GRPO używane w modelach DeepSeek, sprawiły, że RL przestał być domeną wyłącznie robotyki i gier. Wchodzi do mainstreamu AI. I rośnie zapotrzebowanie na ludzi, którzy rozumieją go od środka, a nie tylko klikają w API.

Ten tutorial jest właśnie odpowiedzią na to zapotrzebowanie. Solidna, techniczna robota – nawet jeśli wymaga od czytelnika pewnego wkładu własnego.