How to run a Monte Carlo Tree Search MCTS for stochastic environment?
I am having trouble understanding how to keep track of the expansion, do I expand all stochastic possibilities and weight the return via their chance of happening?

This is indeed one option you can take. This would be very similar in spirit to the idea of &quot;Expectimax&quot; as a variant of minimax for non-deterministic games, in the sense that you'll include explicit &quot;chance nodes&quot; in your tree. When running into such a chance node later on again during a Selection phase, of a later MCTS iteration, you can just select a path of the tree to follow based on a &quot;dice roll&quot;. Importantly, note that this option is only actually available if you have explicit knowledge of exactly when chance events occur, which states they can lead to, and with which probabilities they lead to different states. We also assume that this is feasible, i.e. that you don't have a crazy high (or infinite) number of slightly different game states you could reach.
An alternative option is to use an &quot;open-loop&quot; variant of MCTS. Your nodes would no longer represent game states, but only be representative of the sequence of actions leading to them. You would no longer store any game states in any nodes, but always regenerate them from scratch when traversing the tree, starting from the root node. You would no longer have any explicit chance nodes, but instead have states being representative of larger sets of states that could possibly be reached by following the corresponding path from the root node. For more on this, see my answer to this other question. The advantage of this approach is that it does not require explicit knowledge of all the possible states you can reach due to chance events, do not need explicit knowledge of the probabilities, and can just sample instead of explicitly enumerating every possible outcome.