DreamerV3 and Muzero

Both DreamerV3 and Muzero are model-based RL algorithms. This article dives deep into the details trying to understand these algorithms and run them on RL environments. For each algorithm, we start from understanding the key components, input, output and loss functions. Next, we look at the training details such as code, train batch size, replay buffer size, learning rate etc. Finally, we train the algorithm on RL environments.

DreamerV3

Overall: DreamerV3 is a model-based RL algorithm which is competitive on a wide range of RL tasks across continuous/discrete actions, visual and low-dimensional inputs with fixed hyper-parameters.

The key components

DreamerV3 contained three main components: world model, actor and critic. The world model is responsible to model the hidden transition dynamic, immediate reward and continuation flag (whether episode terminates given the current state and action). The actor and critic, as usual, are responsible to generate action given state (policy) and estimate the value of states (value function).

World Model Components of DreamerV3 from Paper

The world model consists of encoder, decoder, sequence model, dynamic predictor, reward predictor and continue predictor:

The sequence model maintains a hidden state of the current step h_t.

The encoder encodes the input x_t in combination of the hidden state h_t into a latent embedding z_t.

The decoder reconstructs the input x_t from the hidden state h_t and latent embedding z_t.

The dynamic predictor predicts the current latent embedding z_t given the hidden state h_t.

The sequence model generates the next hidden state h_t+1 given the current hidden state h_t, latent embedding z_t and current action a_t.

The reward predictor predicts the immediate reward given z_t and h_t.

The continue predictor predicts the episode continuation flag given z_t and h_t.

Note that the reward is predicted without action a_t.

Components of world model are trained jointly by the weighted sum of three loss terms: prediction loss, dynamic loss and representation loss.

The prediction loss trains the decoder, reward predictor and continue predictor.

The dynamic loss trains dynamic predictor to fit the encoder output

The representation loss trains the encoder to fir the dynamic predictor output.

Why sequence model is not trained?

Sequence model is actually a recurrent process like Gated Recurrent Unit. Therefore, no need to train

The data used to train world model is sampled from replay buffer. The replay buffer store real environment interactions in which the action is sampled from the actor network output (action distribution given a state)

Now, we move to actor and critic learning. It is important to note that both actor and critic are trained from imaged data which is generated by actor, sequence model, dynamic model, reward predictor and continue predictor jointly.

Actor and Critic from Paper

To elaborate this more clearly, suppose we have a replay buffer storing some real environment inputs x. The imaged data is generated in the following way:

1. Sample an input from the replay buffer called x_t.
2. The encoder and sequence model convert x_t into z_t and h_t
3. The actor generate a_t given z_t and h_t
4. The dynamic predictor generate next state z_t+1
5. The rollout continues until horizon=15

Given the rollout imaged data,

The Critic is trained to estimate the expected state value.

Critic loss function from Paper

The actor is trained to maximize the return as well as maintain sufficient exploration via a policy entropy regularizer.

Actor loss function from Paper

Now, we finish the main components of DreamerV3. However, there are many tricks proposed in the paper which contributes largely to DreamerV3’s excellent performance. Without details, we list these tricks below:

1. Symlog function is applied to reward function, encoder input, decoder target and state value target. Symlog function compress both the large positive and negative value, while preserves the sign.

Symlog Function from Paper

2. The target of critic is the boostrapped lambda-returns

Boostrapped lambda-returns from Paper

3. Discrete regression approach for learning the critic based on twohot encoded targets. Returns are transformed using the symlog function and discretize the resulting range into a sequence B of K = 255 equally spaced buckets. The critic network outputs a softmax distribution over the buckets and its output is formed as the expected bucket value under this distribution.

Return target from Paper

4. When train actor, DreamerV3 propose to scale down large returns without scaling up small returns. Scale returns by an exponentially decaying average of the range from their 5th to their 95th batch percentile

Scale down returns from paper

DreamerV3 performance excellent across various tasks as shown in the following figure

Results of DreamerV3 from paper

Training Details

1. The encoder and decoder use convolutional neural networks (CNN) for visual inputs and multi-layer perceptrons (MLPs) for low-dimensional inputs.
2. The dynamics, reward, and continue predictors are also MLPs.
3. Replay buffer size = 1e6, batch_size=16, batch_length=64
4. Activation = LayerNorm + SiLU
5. World model: loss weights = 1.0, 0.5, 0.1; learning_rate =1e-4
6. Actor and critic: learning_rate=3e-5, lambda=0.95, imagation_horizon=15

On training ratio

As we can see, the world model and actor/critic are trained separately. The world model is trained on real environment interaction (replay buffer), while actor and critic are trained on imaged data. How training frequency ratio is an important factor to consider in practice.

training_ratio: The ratio of total steps trained (sum of the sizes of all batches ever sampled from the replay buffer) over the total env steps taken (in the actual environment, not the dreamed one). For example, if the training_ratio is 1024 and the batch size is 1024, we would take 1 env step for every training update: 1024 / 1. If the training ratio is 512 and the batch size is 1024, we would take 2 env steps and then perform a single training update (on a 1024 batch): 1024 / 2.

— — from Ray RLlib DreamerV3

Code

We can find several opens-sourced DreamerV3 algorithms:

1. DreamerV3 Tensorflow and Jax
2. DreamerV3 Pytorch
3. DreamerV3 Tensorflow Ray RLlib

MuZero

Muzero is a model-based RL algorithm equipped with MCTS. Muzero builds on AlphaZero’s powerful search and policy iteration algorithms, but incorporates a learned model into to the training procedure. Muzero achieves state-of-the-art performance ion 57 Atari games and matched the superhuman performance of the AlphaZero.

World Model produces policy, value function, transition dynamic and reward

Representation model encodes historical observations into a hidden state

The dynamic model predict the immediate state and next hidden state given the current state and action

The prediction model generates the policy and state value given the hidden state

All models are trained jointly

The policy loss trains the prediction model to generate the policy close to the search policy (action selection policy of MCTS)

The value loss trains the prediction model to accurately estimate the value target from MCTS and real environment interactions

The reward loss trains the dynamic model to accurately predict the observed immediate reward

L2 regularization is added on the model parameters

MCTS with upper confidence bound

Muzero acts in the environment by selecting action from the MCTS search policy. The MCTS search policy is obtained by running imaged simulations over the learned model

At each real step, a number of MCTS simulations are conducted over the learned model: give the current state, the hidden state is obtained from representation model, an action is selected according to MCTS node statistics. The next hidden state and reward is predicted by the dynamic model and reward model. The simulation continues until a leaf node is reaches. New node is expanded. The node statistics along the simulated trajectory is updated.

How to know whether a leaf node is researched, as we are using the learned model to rollout?

Replay Buffer

For each step, the action is selected from MCTS policy. The environment receives the action and generates new observation and reward. At the end of each episode, the trajectory is stored into the replay buffer.

Train The Model

A trajectory is sampled from the replay buffer. For the initial step, the representation model generates the initial hidden state. Next, the model unroll recurrently for K steps staring from the initial hidden state. At each unroll step k, the dynamic model takes into hidden state and actual action (from the sampled trajectory) and generates next hidden state and reward. The prediction model generated policy and reward. Finally, models are trained with their corresponding target and loss terms defined above.

Code

We can find open-sourced Muzero algorithm

1. Muzero-general pytorch
2. LightZero

Difference between DreamerV3 and Muzero

1. Both learns world model: representation (encoder), dynamic model (recurrent process), reward model on real environment interaction data.
2. Both learns actor and critic (policy model and value model)
3. DreamerV3 trains policy to find action that maximize imaged state value, while Muzero trains policy model to fit MCTS search policy
4. DreamerV3 trains both actor and critic purely on imaged rollout data
5. Muzero trains value and reward model on real environment interaction data.
6. Muzero employs MCTS simulations for each step action slection, while DreamerV3 does not.

In next article, we will run DreamerV3 and Muzero in action.