Torax reinforcement learning environment

[YannBerthelot]

Hi,

This might be a bit off-topic, but I was wondering if there are any plans to make Torax available as a reinforcement learning (RL) environment?

I’m not very familiar with tokamak simulation or control, so apologies in advance if this doesn’t make sense. I’ve seen that DeepMind has applied RL to tokamak control (using a different simulator, I think), and it made me wonder if something similar could be done with Torax.

Having an RL-compatible interface could make it easier to benchmark algorithms on a real physics-based task, which is something the RL community is often missing. It could also help encourage more work on algorithms designed for realistic physical systems.

Thanks a lot for your time — and for the great work on Torax!

[Nush395]

Hi @YannBerthelot,

Thanks for the question, and not at all, it's a great question! Yes previously the work was using RL to do magnetic control.

We are looking to expose a more advanced control-oriented API in the near future which should be quite useful for a RL use case. Part of this API would likely be exposed through the existing SimulationStepFn object https://github.com/google-deepmind/torax/blob/main/torax/_src/orchestration/step_function.py#L67.

This currently has a __call__ method that pushes TORAX forwards one time step (which could be of variable size dependent on if sawteeth are encountered or an adaptive step needs to be taken) and a step method which pushes the simulation forward by a given time step but does not do adaptive stepping (so may not be a converged solution) or sawteeth so the aim is to consolidate this and expose as part of a public stable API.

We're still iterating on the exact interface but the advanced API would expose this object with a method that can take a step for a given time interval with complete physics as configured and this step would be useful in a RL environment.

Will share an interface with you to give a better idea of that is helpful!

[YannBerthelot]

Thank you for your answer ! That is amazing news, this interface would be great indeed.

Sorry if I over explain in the following paragraph, as I do not know how much knowledge you have about RL, this is mainly a way to improve clarity for potential other readers. To be RL-compatible, we need:

• A reset function that (re)sets the system in some initial state (with some degree of stochasticity usually, maybe that does not apply here? I would assume it does, as you cannot exactly control the state of the system ?)
• A step function that takes an action and returns the new state of the system, the SimulationStepFn does seem like a very good start. I may misunderstand what you said but it seems that currently there is no notion of "action", the simulation is 'just' run for a given time. I would assume that the action/control is implicitly communicated through the constraints on the system? If so, do you think it would be possible to vary some of these constraints to transform them into actions for a RL agent?
• A reward function (usually implemented as part of step but conceptually different), that takes the state of the system after an action and returns some numerical signal that specifies how desirable the new state is. Is this something that would be easy to design for Torax?

I do not know how different Torax and the simulator used by Deepmind for Magnetic control of tokamak plasmas, but most likely the same design choices can be transposed to Torax in order to make it RL-compatible.

If step and reset were exposed, and some formula for the reward was designed I would gladly implement the Torax RL environment (while maintaining end-to-end Jax compatibility). If you think that's something in the realm of possibility, I would love to collaborate with you on that !

[Nush395]

Cool!

For the various points you raised above:

• the step function is a stateless JAX function so I'm not sure any state resetting would need to be done? There are functions in TORAX atm for retrieving initial state for a given configuration https://github.com/google-deepmind/torax/blob/main/torax/_src/orchestration/initial_state.py#L38
• that's a good question. The notion of what an action is, isn't something that we want to define in TORAX atm as it could be quite problem dependent but as part of the advanced API we initially plan to expose overrides for a given time steps' runtime parameters, depending on the RL problem at hand these could be thought of as "actions". Not exactly sure what you mean by constraints on the system here though but hopefully that provides clarity.
• Again a reward function isn't something we'd want to define in TORAX atm as it could be quite problem dependent and reward specification is complex in itself but definitely something could compute a reward from the outputs of the TORAX step function.

In general our roadmap atm is to finish making TORAX a physics complete core transport simulator. It would be great to see a RL environment built around TORAX though and we're definitely happy to support with questions and issues in development of that but I think the TORAX library's core responsibility in the near term will be to act as a simulator with applications built on top of it.

Hope that makes sense!

I'll get an initial cut of what the advanced API interface will look like which hopefully will also bring some clarity on how the interface could be used for RL.

[YannBerthelot]

Thanks for the additional info !

• By reset we mean a function that samples an initial state and sets the environment in that state. As you mentioned it is already stateless so the setting part is trivial, but the distribution of initial states is still something that needs to be designed (and that requires knowledge about the system).
• That does provide a lot of clarity, torax is not meant to represent a single task but to allow to construct a lot of different tasks if I understand correctly.
• Similar here it seems.

That makes a lot of sense, I understand your current goal and that designing such a task would be time consuming. I would gladly take a look at the advanced API interface and if it would be easy to use for RL :)

[Nush395]

Just tagging @theo-brown for any additional thoughts/colour as I believe it's something they're interested in.

[theo-brown]

Thanks Nush for the tag! Will chip in some thoughts:

• Anushan has basically covered this, but I think it's worth remembering that fusion is a whole field of science, so there are actually a million different RL problems that could be tackled with the same underlying physics simulator! In this way, your question is analogous to asking researchers in fluid dynamics to make an RL environment for their CFD model. RL is about solving specific tasks (e.g. "make this robot stand up") so in order to design an RL environment in the OpenAI Gym / Farama Gymnasium sense, you actually need to start with a specific task or question (e.g. "get the plasma current from 2MA to 20MA"). Understanding what this means for actuators/actions, constraints, and rewards, is part of the process of building the environment, and will require getting to grips with the underlying field of fusion research - what are people doing, what do they want to do, what obstacles are faced, etc. A Google Scholar search for reinforcement learning tokamak provides an insight into the wide variety of tasks: controlling plasma beta, shape, ramp-down, specific instabilities, etc etc.

• One of our goals at UKAEA is to work towards a set of public benchmark cases for ML problems on the STEP tokamak. There are some developments on this front soon to be announced, so watch this space...

• Your point about stochasticity is a key one, and one that I am aiming to contribute to TORAX in the coming months. Stochastic simulation, while a core part of climate modelling, has yet to be widely adopted in fusion science.

• If you're interested in trying out new ML based design/optimisation algorithms on a fusion problem, you could have a look at the ConStellaration challenge on HuggingFace.

[YannBerthelot]

Thank you for your reply.

• I understand better the aim of Torax, I had the misguided idea that this was somehow an extension of the Deepmind paper on reinforcement learning applied to tokamak control. Now I realize that it's much more complex than this. I was aware that building a new task would require important knowledge about fusion, but I was hoping that the task that was previously addressed by deepmind would be available. My goal was not to try and depict fusion as a simple and single task, sorry if I came across that way.
• I will watch closely ! Once again, if at some point building an RL environment out of any Tokamak tasks becomes interesting and possible to you, I would gladly help.
• That's an interesting point about stochasticity, are the dynamics stochastic in fusion ? or is it only the initial state?
• Sadly, I am currently focused solely on Reinforcement Learning (as I am finishing my PhD in this field), so I do not have much time to try things outside of control tasks. I will take a look at the challenge still, it's very nice to see so much activity around fusion !