def muzero(config: MuZeroConfig):
"""
MuZero training is split into two independent parts: Network training and
self-play data generation.
These two parts only communicate by transferring the latest networks checkpoint
from the training to the self-play, and the finished games from the self-play
to the training.
In contrast to the original MuZero algorithm this version doesn't works with
multiple threads, therefore the training and self-play is done alternately.
"""
storage = SharedStorage(config.new_network(), config.uniform_network(),
config.new_optimizer())
replay_buffer = ReplayBuffer(config)
for loop in range(config.nb_training_loop):
print("Training loop", loop)
score_train = run_selfplay(config, storage, replay_buffer,
config.nb_episodes)
train_network(config, storage, replay_buffer, config.nb_epochs)
print("Train score:", score_train)
print("Eval score:", run_eval(config, storage, 50))
print(
f"MuZero played {config.nb_episodes * (loop + 1)} "
f"episodes and trained for {config.nb_epochs * (loop + 1)} epochs.\n"
)
return storage.latest_network()
def train_network(config: MuZeroConfig, storage: SharedStorage, replay_buffer: ReplayBuffer, epochs: int):
network = storage.current_network
optimizer = storage.current_network
optimizer = storage.optimizer
for _ in range(epochs):
batch = replay_buffer.sample_batch(config.num_unroll_steps, config.td_steps)
update_weights(optimizer, network, batch)
storage.save_network(network.get_training_steps(), network)
def run_eval(config: MuZeroConfig, storage: SharedStorage, eval_episodes: int):
"""Evaluate MuZero without noise added to the prior of the root and without softmax action selection"""
network = storage.latest_network()
returns = []
for _ in range(eval_episodes):
game = play_game(config, network, train=False)
returns.append(sum(game.rewards))
return sum(returns) / eval_episodes if eval_episodes else 0
def run_selfplay(config: MuZeroConfig, storage: SharedStorage, replay_buffer: ReplayBuffer, train_episodes: int):
"""Take the latest network, produces multiple games and save them in the shared replay buffer"""
network = storage.latest_network()
returns = []
for _ in range(train_episodes):
game = play_game(config, network)
replay_buffer.save_game(game)
returns.append(sum(game.rewards))
return sum(returns) / train_episodes