Googles Framework SEED RL soll das ML-Modelltraining in der Cloud beschleunigen

Google-Forscher haben mit SEED RL ein neues Framework für Reinforcement Learning (RL) vorgestellt. Damit Deep-RL-Modelle spezifische Aufgaben sinnvoll lösen, ist bisher ein hoher Trainingsaufwand auf leistungsstarker Infrastruktur erforderlich, der insbesondere in der Cloud auch hohe Kosten verursacht. Mit einem neuen Architekturansatz will das Google-Forscherteam bei SEED RL sowohl die Skalierbarkeit als auch die Trainingseffizienz steigern. Dazu kombiniert das neue Framework zentralisierte Modellinferenz mit einer schnellen Kommunikationsinfrastruktur, die über etablierte RL-Agent-Konzepte wie IMPALA hinausgeht.

Zentralisierte Inferenz

"Mit SEED RL stellen wir einen RL-Agent vor, der auf Tausende von Maschinen skaliert, das Training mit Millionen von Bildern pro Sekunde ermöglicht und die Effizienz der Berechnungen erheblich verbessert", erklärt Lasse Espeholt vom Google-Forscherteam in einem Blogbeitrag zur Vorstellung des neuen RL-Konzepts. Die Inferenz erfolgt dabei zentral auf spezialisierter Hardware mit GPUs oder TPUs. Indem sichergestellt ist, dass die Modellparameter und der Zustand lokal gehalten werden, sollen nicht nur Engpässe bei der Datenübertragung vermeidet, sondern die Inferenz insgesamt beschleunigt werden.

Darüber hinaus soll eine effiziente Netzwerkbibliothek, die auf dem gRPC-Framework mit asynchronen Streaming-RPCs basiert, gewährleisten, dass bei der Übertragung der Beobachtungen an den Learner nach jedem Schritt nur geringe Latenzen auftreten. Während der Learner sich dabei auf Tausende Cores – im konkreten Fall von Cloud-TPUs bis zu 2048 Kerne – skalieren lässt, kann die Zahl der Actors auf Tausende Maschinen verteilt werden.

V-trace- und R2D2-Algorithmen

Damit dieser Architekturansatz funktioniert, integriert das auf der TensorFlow-2-API basierte SEED RL den Google-Forschern zufolge zwei entscheidende Algorithmen. Einerseits V-trace, eine bereits mit IMPALA eingeführte Policy-Gradienten-Methode, und den Q-Learning-Algorithmus R2D2. Um die Leistungsfähigkeit von SEED RL demonstrieren zu können, greift das Team auf gängige RL-Benchmarks wie das Arcade Learning Environment, DeepMind Lab und Google Research Football zurück. Bei den Versuchen in der DeepMind-Lab-Umgebung skalierte SEED RL mit 64 Cloud-TPU-Cores auf bis zu 2,4 Million Bilder pro Sekunde, was dem 80-fachen des bisher mit IMPALA erzielbaren entspricht. Im Google-Research-Football-Benchmark gelang es den Forschern offenbar erstmals die Aufgabe "Hard" zu lösen, in dem sie SEED RL mit einem vergrößerten Modell und höherer Eingabeauflösung nutzten.

Weitere Benchmark-Ergebnisse sowie nähere Details zu SEED RL finden sich im Blogbeitrag des Google-Teams sowie in einem ergänzenden Forschungspapier. Der Quellcode von SEED RL steht als Open Source frei auf GitHub zur Verfügung – einschließlich verschiedener Beispiele, die auf den Einsatz in der Google Cloud mit GPUs ausgelegt sind. (map)