Gym-TORAX: Open-source software for integrating RL with plasma control simulators

Cet article présente Gym-TORAX, un logiciel open-source en Python qui permet de créer des environnements d'apprentissage par renforcement compatibles avec les bibliothèques standard pour simuler et optimiser le contrôle des plasmas dans les tokamaks, en s'appuyant sur le simulateur TORAX et en offrant dès à présent un scénario de démarrage pour ITER.

L'essentiel

🍳 Le Grand Défi : Cuisiner un Soleil dans une Casserole

Imaginez que vous essayez de faire cuire un soleil dans votre cuisine (c'est ce qu'on appelle la fusion nucléaire). Pour cela, vous utilisez un appareil géant en forme de beignet appelé un Tokamak.

Le problème ? Le "feu" à l'intérieur (le plasma) est très capricieux. Si vous ne le surveillez pas à la seconde près, il s'éteint ou explose. C'est comme essayer de faire un soufflé parfait pendant un tremblement de terre : il faut ajuster les boutons de température et de courant des milliers de fois par seconde.

Jusqu'à présent, seuls des experts en physique très pointus pouvaient programmer ces ajustements. C'était comme demander à un chef étoilé de cuisiner sans recette, juste en devinant.

🤖 L'Idée Géniale : Apprendre au Robot à Cuisiner

C'est là qu'intervient l'Apprentissage par Renforcement (RL). C'est une méthode où l'on donne un robot (une intelligence artificielle) le contrôle de la cuisine. Le robot essaie, se trompe, apprend de ses erreurs, et finit par devenir un chef incroyable.

Mais pour entraîner ce robot, il faut un simulateur (une cuisine virtuelle) où il peut faire des milliers d'essais sans risquer de faire exploser la vraie centrale.

🧰 La Boîte à Outils : Gym-TORAX

C'est exactement ce que propose ce papier : Gym-TORAX.

Imaginez que TORAX est un super moteur de simulation très rapide (écrit en Python) qui calcule comment le plasma bouge. Mais ce moteur est un peu "brut" : il faut savoir coder en physique nucléaire pour l'utiliser. C'est comme avoir une Ferrari mais ne pas savoir conduire.

Gym-TORAX est le volant et le tableau de bord que l'on ajoute à cette Ferrari.

• Il transforme le simulateur complexe en un jeu vidéo simple et standard (appelé "Gymnasium").
• Il permet à n'importe quel développeur en intelligence artificielle (même sans être physicien) de dire : "Voici ce que je veux observer (la température), voici ce que je peux contrôler (les boutons), et voici mon objectif (garder le soufflé gonflé)."

🎮 Comment ça marche ? (L'analogie du Jeu Vidéo)

Le papier explique que Gym-TORAX crée un environnement de jeu en trois étapes simples :

1. Le Scénario (La Recette) : On définit une situation, par exemple le démarrage d'un réacteur (comme le projet ITER).
2. Les Actions (Les Boutons) : L'IA peut appuyer sur des boutons virtuels pour changer le courant électrique ou ajouter de la chaleur.
3. Les Points (La Récompense) :
   • Si le plasma est stable et performant ➡️ L'IA gagne des points (+100).
   • Si le plasma devient instable ➡️ L'IA perd des points (-1000) et la partie s'arrête (Game Over).

L'IA joue des millions de parties contre elle-même pour trouver la stratégie parfaite pour gagner le plus de points possible.

📈 Ce que les auteurs ont testé

Dans la partie "Annexe" du papier, ils ont fait un petit tour de force :

• Ils ont pris un scénario réel (le démarrage d'ITER).
• Ils ont comparé trois "joueurs" :
  1. Le Robot "Lecteur de Recette" (Open-loop) : Il suit une recette pré-écrite. Ça marche bien, mais pas mieux que prévu.
  2. Le Robot "Détrousseur" (Random) : Il appuie sur les boutons au hasard. Catastrophe ! Le plasma s'effondre.
  3. Le Robot "Chef Expert" (Contrôleur PI) : Un robot un peu plus intelligent qui ajuste les boutons en temps réel. Résultat : Il bat la recette de base et obtient un meilleur résultat !

🚀 Pourquoi c'est important ?

Avant, pour faire de la recherche sur le contrôle des réacteurs à fusion, il fallait être un physicien nucléaire ET un expert en informatique. C'était une barrière énorme.

Gym-TORAX change la donne :

• Il rend le code libre et gratuit (Open-source).
• Il permet aux experts en IA de travailler sur la fusion sans avoir à apprendre 10 ans de physique.
• Il permet aux physiciens de tester de nouvelles idées de contrôle très rapidement.

En résumé : Gym-TORAX est le pont qui permet aux intelligences artificières d'apprendre à piloter les réacteurs du futur, rendant l'énergie propre et illimitée un peu plus proche de notre réalité. C'est comme donner un manuel de conduite simplifié à un pilote de chasse pour qu'il puisse apprendre à piloter un avion de pointe sans passer 20 ans à l'école de l'air.

Résumé technique

Titre du papier

Gym-TORAX : Logiciel open-source pour l'intégration de l'apprentissage par renforcement avec des simulateurs de contrôle de plasma dans la recherche sur les tokamaks.

---

1. Problématique

La recherche sur l'énergie de fusion vise à optimiser la stabilité et les performances des réacteurs à tokamak. Cependant, le contrôle et la conception de ces dispositifs sont extrêmement complexes en raison de :

• La haute dimensionnalité du problème.
• La présence de nombreuses non-linéarités inhérentes à la dynamique du plasma.

Bien que l'Apprentissage par Renforcement (RL) ait montré des promesses dans d'autres domaines (robotique, réseaux électriques), son application au contrôle du plasma est freinée par le manque d'outils accessibles. Les simulateurs existants sont souvent :

• Fermés (licences restrictives, ex: RAPTOR, JOREK).
• Conçus pour les physiciens du plasma, rendant leur utilisation difficile pour les chercheurs en RL.
• Dépourvus d'interfaces adaptées aux applications de contrôle en boucle fermée (ex: EFIT ne fait que de la reconstruction d'équilibre).

2. Méthodologie

L'approche proposée consiste à créer une interface standardisée entre les algorithmes de RL et le simulateur physique TORAX.

A. Architecture Logicielle

• Base physique : Le package s'appuie sur TORAX, un simulateur open-source écrit en Python utilisant JAX pour une différenciation automatique rapide et une exécution performante. TORAX résout des équations aux dérivées partielles (PDE) décrivant l'évolution des températures, densités et flux magnétiques du plasma.
• Interface RL : Gym-TORAX agit comme un wrapper autour de TORAX, implémentant l'API standard Gymnasium. Cela permet d'encapsuler la physique complexe derrière une interface familière pour les chercheurs en RL.

B. Modélisation du Problème de Contrôle

Le problème est formulé comme un Processus de Décision Markovien (MDP) fini et déterministe $(\mathcal{S}, \mathcal{A}, f, r, s_0, \gamma, T)$ :

• Espace d'état ($\mathcal{S}$) : Comprend les variables d'état du plasma (températures, densités, flux) et des métriques dérivées (facteur de sécurité $q$, gain de fusion $Q$, etc.).
• Espace d'action ($\mathcal{A}$) : L'agent contrôle des variables temporelles telles que la tension de boucle ($V_{loop}$), le courant total ($I_p$) et les sources d'énergie (chauffage par injection de neutres NBI, chauffage par résonance cyclotronique électronique ECRH).
• Fonction de transition : Chaque étape de transition de l'agent correspond à la résolution de $K$ pas de temps par TORAX. Deux modes de discrétisation sont disponibles : auto (pas de temps adaptatif) ou fixed.
• Fonction de récompense ($r$) : Définie par l'utilisateur pour optimiser des objectifs spécifiques (stabilité, performance, confinement). Elle pénalise les états non viables (arrêt de la simulation avec une récompense négative élevée).

C. Création d'Environnements

Pour créer un nouvel environnement, l'utilisateur hérite de la classe BaseEnv et implémente quatre méthodes abstraites :

1. _get_torax_config(): Définit les conditions initiales et le modèle physique.
2. _define_action_space(): Spécifie les actions contrôlables et leurs limites de taux de variation (ramp-rate).
3. _define_observation_space(): Sélectionne les variables observables par l'agent.
4. _compute_reward(): Définit la fonction de récompense.

3. Contributions Clés

• Premier framework open-source dédié à l'intégration directe de TORAX avec les bibliothèques RL via l'API Gymnasium.
• Réduction de la barrière d'entrée : Permet aux chercheurs en RL de se concentrer sur les stratégies de contrôle sans maîtriser la physique des plasmas complexe.
• Environnement de référence prêt à l'emploi : Inclut une implémentation complète basée sur le scénario de "ramp-up" hybride du réacteur ITER (International Thermonuclear Experimental Reactor).
• Facilitation de la collaboration : Crée un pont entre la communauté de la physique des plasmas et celle de l'intelligence artificielle.

4. Résultats (Étude de Cas)

Les auteurs ont évalué trois politiques dans l'environnement IterHybridEnv (scénario de montée en puissance de 150 secondes, passant du mode L au mode H) :

1. Politique en boucle ouverte ($\pi_{OL}$) : Suit les trajectoires prédéfinies du scénario de référence.
   • Retour espéré : 3.40
2. Politique aléatoire ($\pi_{R}$) : Sélectionne les actions au hasard.
   • Retour espéré : -10.79 (échec dû à l'instabilité du plasma).
3. Contrôleur PI optimisé ($\pi_{PI}$) : Utilise un contrôleur Proportionnel-Intégral pour réguler le courant central, avec des gains ($k_p, k_i$) optimisés par recherche sur grille.
   • Retour espéré : 3.79

Analyse des résultats :

• La politique PI surpasse le scénario de référence (amélioration de ~11 % du retour).
• La trajectoire du courant générée par la politique PI montre une augmentation stable jusqu'à la limite maximale autorisée (15 MA), ce qui est corrélé à un meilleur confinement.
• Les résultats démontrent que même des contrôleurs classiques optimisés peuvent servir de base solide pour des algorithmes RL plus avancés.

5. Signification et Perspectives

Signification :
Gym-TORAX est un outil catalyseur pour la recherche sur le contrôle des tokamaks. Il transforme un problème physique complexe en un problème d'optimisation standardisable, permettant d'explorer de nouvelles stratégies de contrôle qui seraient difficiles à concevoir manuellement. Sa nature open-source favorise la reproductibilité et la communauté.

Perspectives Futures :

• Paramétrisation géométrique : Permettre de définir la géométrie du tokamak et du plasma directement lors de la création de l'environnement (ajout d'une nouvelle dimension au problème RL).
• Gestion d'événements physiques : Intégration d'outils pour gérer des événements critiques comme la transition L-H (changement de régime de confinement).
• Extension des capacités : Profiter des évolutions futures du simulateur TORAX pour enrichir les capacités de Gym-TORAX.

En conclusion, ce logiciel ouvre la voie à l'application de l'apprentissage par renforcement profond pour le contrôle en temps réel des réacteurs à fusion de demain.