xBerry Robotic Arm RL Tydzień 3: Trening Fazy B z kamerą i architekturą A2C

Po tygodniach treningu na idealnych danych numerycznych zabraliśmy robotowi ściągawkę. Od tej chwili musi uczyć się z tego, co widzi.

TL;DR

Tydzień 3 to przejście z Fazy A (tylko propriocepcja) do Fazy B: trening na surowych klatkach z kamery. Wykorzystaliśmy wagi Krytyka z Fazy A jako warm start, zbudowaliśmy asymetryczną architekturę A2C, w której Aktor widzi tylko piksele i pozycje przegubów, a Krytyk zachowuje dostęp do uprzywilejowanych danych i osiągnęliśmy 96% skuteczności (48/50 prób) przy przenoszeniu piłki w symulacji.

Od liczb do pikseli

Faza A dała nam solidne podstawy: politykę wytrenowaną na idealnych danych numerycznych — dokładnych pozycjach przegubów, współrzędnych piłki i miski – działającą równolegle w ponad 4000 środowiskach symulacyjnych. Ta polityka osiągnęła niemal doskonałą skuteczność w symulacji. Był jednak jeden problem: na prawdziwym robocie większość tych uprzywilejowanych danych po prostu nie istnieje. Ramię SO-101 nie otrzymuje podczas inferencji dokładnych współrzędnych piłki. Widzi to, co widzą jego kamery.

Faza B zamyka tę lukę. Celem jest wytrenowanie polityki, która potrafi działać tylko na podstawie klatek z kamery – tego samego strumienia informacji dostępnego podczas wdrożenia w rzeczywistym świecie.

Trik z warm startem

Przejście z liczb na piksele brzmi jak zaczynanie od nowa. Tak nie jest i ta różnica ma ogromne znaczenie dla czasu treningu.

Przenieśliśmy wagi Krytyka bezpośrednio z Fazy A do Fazy B. Krytyk – sieć neuronowa odpowiedzialna za ocenę jakości podejmowanych akcji – już wie, jak wygląda dobra praca ramienia. Spędził całą Fazę A ucząc się oceniać akcje na podstawie bogatych danych stanowych. Dzięki inicjalizacji Fazy B tymi wagami Aktor (właściwy sterownik ruchu) otrzymuje wiarygodne sygnały zwrotne już od pierwszego kroku treningowego, zamiast uczyć się od Krytyka, który sam jeszcze szuka właściwej drogi.

Dlaczego to ważne: Trening polityk wizyjnych w RL jest powolny i kosztowny obliczeniowo. Bez warm startu Krytyk potrzebowałby tysięcy kroków, zanim jego oceny akcji stałyby się wartościowe. Warm start kompresuje ten koszt kalibracji do zera.

Asymetryczne A2C: jedna polityka, dwie różne rzeczywistości

Centralnym wyzwaniem architektonicznym Fazy B jest to, że Aktor i Krytyk muszą otrzymywać fundamentalnie różne rodzaje informacji.

A2C (Advantage Actor-Critic) – algorytm uczenia ze wzmocnieniem, w którym dwie sieci współpracują: Aktor decyduje, jaką akcję wykonać, a Krytyk ocenia, na ile była ona dobra.

Podczas inferencji na rzeczywistym robocie działa tylko Aktor. Krytyk trenuje razem z nim, ale nigdy nie jest wdrażany. Ta asymetria stwarza szansę: Krytyk może otrzymywać uprzywilejowane informacje – dokładną pozycję piłki i miski wyrażoną jako kwaternion (x, y, z, w) – niedostępne w rzeczywistym świecie, podczas gdy Aktor uczy się wyłącznie z tego, do czego będzie miał dostęp prawdziwy robot.

Aktor otrzymuje: osadzenia z kamery (skompresowane reprezentacje wizualne z enkodera CNN) oraz pozycje przegubów z enkoderów serw robota.

Krytyk otrzymuje wszystko, co widzi Aktor, plus uprzywilejowane obserwacje dotyczące dokładnych pozycji 3D piłki i miski.

Jak nauczyliśmy robota widzieć: enkoder CNN

Surowe klatki z kamery nie mogą bezpośrednio trafiać do sieci polityki. Pojedynczy obraz RGB o rozdzielczości 480×640 zawiera setki tysięcy wartości pikseli – podawanie tego bez przetworzenia do sieci neuronowej jest obliczeniowo niepraktyczne i statystycznie zaszumione.

CNN (Convolutional Neural Network) — typ sieci neuronowej zaprojektowany do przetwarzania danych przestrzennych, takich jak obrazy, wykrywający cechy przez wyuczone filtry stosowane do całego obrazu.

Nasz PixelEncoder składa się z 4 warstw konwolucyjnych, po których następuje SpatialSoftmax — warstwa obliczająca oczekiwaną pozycję 2D każdej wyuczonej cechy w obrazie, skutecznie zamieniająca mapę cech w zestaw współrzędnych przestrzennych. Końcowa głowica liniowa przekształca te współrzędne w kompaktowy wektor osadzeń.

Osadzenia (embeddingi) – kompaktowe wektory liczbowe reprezentujące obraz w przestrzeni latentnej. Podobne obrazy generują podobne wektory osadzeń, co pomaga polityce generalizować w różnych warunkach wizualnych.

class SpatialSoftmax(nn.Module):
 def __init__(self):
 super().__init__()
 self._grid_cache: dict[tuple[int, int, torch.device], tuple] = {}
 def _grids(self, h: int, w: int, device: torch.device):
 key = (h, w, device)
 cached = self._grid_cache.get(key)
 if cached is None:
 xs = torch.linspace(-1, 1, w, device=device)
 ys = torch.linspace(-1, 1, h, device=device)
 y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij")
 cached = (x_grid.reshape(-1), y_grid.reshape(-1))
 self._grid_cache[key] = cached
 return cached
 def forward(self, feat):
 B, C, H, W = feat.shape
 x_grid, y_grid = self._grids(H, W, feat.device)
 w = torch.softmax(feat.reshape(B, C, H * W), dim=-1)
 x = (w * x_grid).sum(-1)
 y = (w * y_grid).sum(-1)
 return torch.cat([x, y], dim=-1)
class PixelEncoder(nn.Module):
 def __init__(self, in_channels=3, feature_dim=128):
 super().__init__()
 self.conv = nn.Sequential(
 nn.Conv2d(in_channels, 32, 3, stride=2),
 nn.ReLU(),
 nn.Conv2d(32, 32, 3, stride=1),
 nn.ReLU(),
 nn.Conv2d(32, 32, 3, stride=1),
 nn.ReLU(),
 nn.Conv2d(32, 32, 3, stride=1),
 nn.ReLU(),
  )
 self.spatial_softmax = SpatialSoftmax()
 self.head = nn.Sequential(
 nn.Linear(2 * 32, feature_dim),
 nn.LayerNorm(feature_dim),
 nn.Tanh(),
  )
 def forward(self, x):
 return self.head(self.spatial_softmax(self.conv(x)))

Efekt: podczas inferencji surowe klatki z kamery są zamieniane na wektor o stałym rozmiarze, który sieć Aktora potrafi interpretować. Kodowanie dodaje minimalną latencję i umożliwia działanie całej polityki wyłącznie na pikselach.

Sufit pamięci

Przejście na trening wizyjny ujawniło twardy fizyczny ogranicznik: pamięć GPU.

Faza A komfortowo działała z 4096 równoległymi środowiskami symulacyjnymi jednocześnie. Obserwacje kamerowe w Fazie B wymagają przechowywania tensorów obrazu dla każdego środowiska w każdym kroku treningowym. Przy 4096 środowiskach łączna pamięć dla wszystkich tensorów obrazu niemal natychmiast wyczerpuje GPU. Rozwiązaniem była metoda prób i błędów: systematyczne zmniejszanie liczby środowisk, aż trening działa bez błędów OOM (out-of-memory). Wylądowaliśmy na około 64 równoległych środowiskach – redukcja 64× w stosunku do Fazy A.

To znacznie spowalnia przepustowość na krok, ale warm startowany Krytyk rekompensuje to, sprawiając, że każdy krok gradientowy przynosi więcej wartości od samego początku.

Wyniki: 48 na 50

Po wytrenowaniu Fazy B z asymetryczną architekturą A2C opartą na obrazach przeprowadziliśmy ewaluację na 50 próbach. Polityka odniosła sukces w 48 z nich – 96% skuteczności na zadaniu przenoszenia piłki.

Co istotne, wynik jest porównywalny z Fazą A, mimo że polityka podczas ewaluacji otrzymuje tylko klatki z kamery i pozycje przegubów. Warm startowany Krytyk, dysponujący uprzywilejowanymi danymi stanowymi podczas treningu, dał Aktorowi wystarczający sygnał do wykształcenia solidnej polityki wizyjnej – bez potrzeby dostępu do idealnych danych stanowych w czasie inferencji.

Co dalej?

Tydzień 3 kończy się checkpointem zwalidowanym w symulacji z 96% skutecznością. Tydzień 4 to wdrożenie: model trafia na fizyczne ramię SO-101. Prawdziwe pytanie nie brzmi już, czy polityka działa w symulacji – brzmi, czy cechy wizualne, których się nauczyła, generalizują się na prawdziwe oświetlenie, prawdziwe kamery i prawdziwą fizykę.

FAQ

Czym jest Faza B w treningu robotów metodą uczenia ze wzmocnieniem?

Faza B to etap treningowy, w którym polityka robota uczy się działać na podstawie surowych obrazów z kamery zamiast uprzywilejowanych danych numerycznych, bezpośrednio przygotowując politykę do wdrożenia w rzeczywistym świecie, gdzie dokładne pozycje obiektów są niedostępne.

Czym jest asymetryczna architektura A2C?

Asymetryczne A2C (Advantage Actor-Critic) to układ sieci neuronowych, w którym Aktor otrzymuje tylko obserwacje dostępne podczas rzeczywistego wdrożenia (osadzenia z kamery, pozycje przegubów), podczas gdy Krytyk dodatkowo otrzymuje uprzywilejowane informacje (dokładne współrzędne obiektów) w celu zapewnienia dokładniejszego sygnału treningowego.

Czym jest SpatialSoftmax i dlaczego jest używany w uczeniu robotów?

SpatialSoftmax zamienia mapę cech konwolucyjnych w oczekiwane współrzędne 2D dla każdego kanału cech, skutecznie lokalizując kluczowe elementy wizualne w obrazie. W uczeniu robotów pomaga enkodowi wykrywać i śledzić obiekty istotne dla zadania – takie jak piłka czy miska – w kolejnych klatkach z kamery.

Czym jest warm start w treningu RL?

Warm start oznacza inicjalizację nowego procesu treningowego wagami z poprzednio wytrenowanej sieci, zamiast losowej inicjalizacji. W Fazie B wagi Krytyka z Fazy A są ponownie wykorzystywane, dzięki czemu Aktor otrzymuje sensowne oceny akcji natychmiast, znacznie skracając całkowity czas treningu.

Dlaczego trening na obrazach z kamery zmniejsza liczbę równoległych środowisk?

Każda obserwacja kamerowa to tensor obrazu zawierający setki tysięcy wartości dla każdego środowiska. Przy 4096 równoległych środowiskach łączna pamięć dla wszystkich tensorów obrazu przekracza pojemność GPU. Redukcja do około 64 środowisk utrzymuje zużycie pamięci w bezpiecznych granicach, zachowując stabilność treningu.

Jaki wynik osiągnął robot podczas ewaluacji Fazy B w symulacji?

Robot osiągnął 96% skuteczności – 48 udanych prób przeniesienia piłki na 50 – używając podczas ewaluacji wyłącznie klatek z kamery i pozycji przegubów jako wejścia, dorównując wynikowi polityki wytrenowanej na propriocepcji w Fazie A.