How to change muzero to allow the invertible transformation

[ipsec]

Hi all,

First, congratulations by excellent project.

On the muzero paper (https://rdcu.be/ccErB) in the Network Architecture they have used an invertible transformation, for reward and value targets, to a categorical representation.

The code of Werner Duvaud available on GitHub had done this here and here.

Using this, the loss function could be the cross entropy given more stable results than MSE (according the paper).

I don't know how to implement this on the alf.
Maybe some guidance, so I can try help you with this.

Best regards,
Fernando

[emailweixu]

Hi, Fernando,

Thanks for offering the help.

You need to change MCTSModel.calc_loss in mcts_mode.py to do this. You can see a recent change (#1158) to use quantile regression for reward and value regression to get an idea where are the relevant code. And if you are using create_simple_prediction_net() (the default to SimpleMCTSModel) to construct the prediction model, you also need to change create_simple_prediction_net() to output the logits instead of a scalar.

I'd love to know how it compares to quantile regression.

Don't hesitate to ask if there is still any confusion.

Wei

[emailweixu]

@ipsec,

FYI, I am implementing this.

[ipsec]

[ipsec]

[ipsec]

Hi @emailweixu,

I saw your working in the PR (https://github.com/HorizonRobotics/alf/compare/PR_muzero_categorical_loss_for_reward_value) and I'm trying to use muzero_atari_conf.py like a template to create a cartpole version, but I didn't have success.

Have you a config version using this PR of Cartpole-v1? So I can test it.

Thanks in advanced.

[emailweixu]

There is a config for Atari: examples/muzero_atari_conf.py

[ipsec]

Hi @emailweixu

The examples/muzero_atari_conf.py is a bit complex to me (I just want to run CartPole-v1 env).
I'm trying to simplify it and I'm getting erros like this:

W0211 10:16:05.978615 4677125632 muzero_cartpole_conf.py:464] The value of config 'TrainerConfig.debug_summaries' has been configured to True. It is replaced by the new value True
W0211 10:16:05.979274 4677125632 muzero_cartpole_conf.py:464] The value of config 'TrainerConfig.summarize_grads_and_vars' has been configured to True. It is replaced by the new value True
I0211 10:16:05.983580 4677125632 parallel_environment.py:94] Spawning all processes.
I0211 10:16:06.124400 4677125632 parallel_environment.py:101] All processes started.
I0211 10:16:06.167684 4677125632 policy_trainer.py:332] observation_spec=BoundedTensorSpec(shape=(4,), dtype=torch.float32, minimum=array([-4.8000002e+00, -3.4028235e+38, -4.1887903e-01, -3.4028235e+38],
      dtype=float32), maximum=array([4.8000002e+00, 3.4028235e+38, 4.1887903e-01, 3.4028235e+38],
      dtype=float32))
I0211 10:16:06.168600 4677125632 policy_trainer.py:334] action_spec=BoundedTensorSpec(shape=(), dtype=torch.int64, minimum=array(0), maximum=array(1))
I0211 10:16:06.173171 4677125632 policy_trainer.py:343] transformed_observation_spec=BoundedTensorSpec(shape=(4,), dtype=torch.float32, minimum=array([-4.8000002e+00, -3.4028235e+38, -4.1887903e-01, -3.4028235e+38],
      dtype=float32), maximum=array([4.8000002e+00, 3.4028235e+38, 4.1887903e-01, 3.4028235e+38],
      dtype=float32))
I0211 10:16:06.700411 4677125632 parallel_environment.py:168] Closing all processes.
I0211 10:16:07.022981 4677125632 parallel_environment.py:171] All processes closed.
Traceback (most recent call last):
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/Users/fernando/PycharmProjects/alf/alf/bin/train.py", line 202, in main
    training_worker(
  File "/Users/fernando/PycharmProjects/alf/alf/bin/train.py", line 177, in training_worker
    raise e
  File "/Users/fernando/PycharmProjects/alf/alf/bin/train.py", line 171, in training_worker
    _train(root_dir, rank, world_size)
  File "/Users/fernando/PycharmProjects/alf/alf/bin/train.py", line 140, in _train
    trainer.train()
  File "/Users/fernando/PycharmProjects/alf/alf/trainers/policy_trainer.py", line 148, in train
    self._restore_checkpoint()
  File "/Users/fernando/PycharmProjects/alf/alf/trainers/policy_trainer.py", line 489, in _restore_checkpoint
    super()._restore_checkpoint(checkpointer)
  File "/Users/fernando/PycharmProjects/alf/alf/trainers/policy_trainer.py", line 286, in _restore_checkpoint
    self._algorithm.train_iter()
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/rl_algorithm.py", line 538, in train_iter
    return self._train_iter_off_policy()
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/rl_algorithm.py", line 569, in _train_iter_off_policy
    experience = self.unroll(config.unroll_length)
  File "/Users/fernando/PycharmProjects/alf/alf/utils/common.py", line 997, in _func
    ret = func(*args, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/utils/distributed.py", line 194, in wrapped
    return method(*args, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/rl_algorithm.py", line 473, in unroll
    policy_step = self.rollout_step(transformed_time_step,
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/rl_algorithm.py", line 426, in _rollout_step
    policy_step = self._original_rollout_step(time_step, state)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/agent.py", line 282, in rollout_step
    rl_step = self._rl_algorithm.rollout_step(rl_time_step, state.rl)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/rl_algorithm.py", line 426, in _rollout_step
    policy_step = self._original_rollout_step(time_step, state)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/muzero_algorithm.py", line 260, in rollout_step
    return self._mcts.predict_step(time_step, state)
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/mcts_algorithm.py", line 428, in predict_step
    model_output = self._model.initial_inference(time_step.observation)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/mcts_models.py", line 224, in initial_inference
    model_output = self._predict(model_state)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/mcts_models.py", line 253, in _predict
    model_output = self.prediction_model(state.state, state.pred_state)
  File "/Users/fernando/PycharmProjects/alf/alf/algorithms/mcts_models.py", line 722, in prediction_model
    game_over_logit), pred_state = self._prediction_net(
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/networks/containers.py", line 406, in forward
    output = list(
  File "/Users/fernando/PycharmProjects/alf/alf/networks/containers.py", line 407, in <lambda>
    map(lambda net: net(inputs)[0], self._networks_flattened))
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/networks/network.py", line 347, in forward
    return self._module(x), state
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/layers.py", line 3122, in _forward_simple
    input = module(input)
  File "/usr/local/Caskroom/miniconda/base/envs/muzero/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/fernando/PycharmProjects/alf/alf/layers.py", line 438, in forward
    y = inputs.matmul(self._weight.t())
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x1024 and 1x1024)

This is the modified muzero_atari_conf.py I'm trying.

# Copyright (c) 2021 Horizon Robotics and ALF Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Configuration for training Atari games using MuZero.

The structure of the big model (use_small_net=False, use_bn=True, lstm_reward=True,
sum_over_reward_prediction=False, predict_reward_sum=True, train_repr_prediction=True)
largely follows the model structure of EfficientZero (arXiv:2111.00210)

There are several important algorithmic differences with EfficientZero:

1. Different loss for reward and value.
    EfficientZero uses:
    - rv_loss = DiscreteRegressionLoss(SqrtLinearTransform(0.001), inverse_after_mean=True)
    We use:
    - rv_loss = OrderedDiscreteRegressionLoss(Sqrt1pTransform(), inverse_after_mean=False)
2. different MCTS.
    EfficientZero uses:
    - num_parallel_sims = 1
    - (act/learn/search)_with_exploration_policy = False
    - root_exploration_fraction = 0.25
    We use:
    - num_parallel_sims = 2
    - (act/learn/search)_with_exploration_policy = True
    - root_exploration_fraction = 0
3. different BatchNorm:
    EfficientZero uses torch.nn.BatchNorm.
    We use customized BatchNorm (alf.layers.BatchNorm) which can correctly handle
    state.
3. different optimizer. EfficientZero uses SGD. We use ADAMTF.

There are many other hyper-parameter differences not listed above.
"""

import torch
import torch.nn as nn
from functools import partial
from alf.environments import alf_wrappers

import alf
# import alf.examples.atari_conf
import alf.examples.muzero_conf
from alf.utils.schedulers import LinearScheduler, StepScheduler
from alf.algorithms.muzero_algorithm import LinearTdStepFunc
from alf.algorithms.mcts_models import SimpleMCTSModel
from alf.algorithms.mcts_algorithm import MCTSAlgorithm, VisitSoftmaxTemperatureByProgress
from alf.optimizers import SGD, Adam, AdamTF, AdamW
from alf.algorithms.data_transformer import RewardClipping

from alf.algorithms.data_transformer import FrameStacker
from alf.utils.summary_utils import summarize_tensor_gradients
from alf.utils import losses

# MuzeroAlgorithm does not suppport the RewardClipping configured in atari_conf
# alf.config('TrainerConfig', data_transformer_ctor=[FrameStacker])


def define_config(name, default_value):
    alf.define_config(name, default_value)
    return alf.get_config_value('_CONFIG._USER.' + name)


train_repr_prediction = define_config('train_repr_prediction', True)
train_game_over_function = define_config('train_game_over_function', False)
num_envs = define_config('num_envs', 1)
unroll_length = define_config('unroll_length', 1)
use_small_net = define_config('use_small_net', True)
discount = define_config('discount', 0.988)
initial_lr = define_config('initial_lr', 2e-4)
norm_type = define_config('norm_type', alf.layers)
optimizer_type = define_config('optimizer_type', 'ADAMTF')
terminal_on_life_loss = define_config('terminal_on_life_loss', False)
use_bn = define_config('use_bn', True)
pred_use_bn = define_config('pred_use_bn', False)
lstm_reward = define_config('lstm_reward', False)
sum_over_reward_prediction = define_config('sum_over_reward_prediction', False)
weight_decay = define_config('weight_decay', 0.1)
rv_weight_decay = define_config('rv_weight_decay', 4e-6)
rv_bias_decay = define_config('rv_bias_decay', 0)
action_weight_decay = define_config('action_weight_decay', None)
action_bias_decay = define_config('action_bias_decay', None)

rv_loss = define_config(
    'rv_loss',
    losses.OrderedDiscreteRegressionLoss(
        transform=alf.math.Sqrt1pTransform(), inverse_after_mean=False))
rv_bias_zero_init = define_config('rv_bias_zero_init', False)

# This option can make the optimization of the parameters of quantile regression
# invariant to the number of quantiles.
scale_grad_by_num_quantiles = define_config('scale_grad_by_num_quantiles',
                                            False)

# alf.math.identity, alf.layers.BatchNorm2d(64), alf.math.normalize_min_max
dyn_state_normalizer = define_config('dyn_state_normalizer', alf.math.identity)

# If 1, not use quantile regression
num_quantiles = define_config('num_quantiles', 256)
alf.config('multi_quantile_huber_loss', delta=0.0)

reward_transformer = RewardClipping() if num_quantiles == 1 else None

alf.config(
    "AverageDiscountedReturnMetric",
    discount=discount,
    reward_transformer=reward_transformer)

# if use_small_net:
#     alf.config('DMAtariPreprocessing', screen_size=84, gray_scale=True)
# else:
#     alf.config('DMAtariPreprocessing', screen_size=96, gray_scale=False)

# if terminal_on_life_loss:
#    alf.config(
#        'suite_gym.load',
#        alf_env_wrappers=[alf_wrappers.AtariTerminalOnLifeLossWrapper])

alf.config(
    "create_environment",
    env_name="CartPole-v1",
    num_parallel_environments=num_envs)

alf.config("alf.batch_norm.BatchNorm1d", affine=True)
alf.config("alf.batch_norm.BatchNorm2d", affine=True)
alf.config(
    "layers.ResidueBlock",
    with_batch_normalization=use_bn,
    bn_ctor=norm_type.BatchNorm2d)
alf.config("layers.Conv2D", use_bn=use_bn, bn_ctor=norm_type.BatchNorm2d)
alf.config("layers.FC", bn_ctor=norm_type.BatchNorm1d)


def create_representation_net(observation_spec):
    in_channels = observation_spec.shape[0]
    return alf.nn.Sequential(
        alf.layers.Scale(1. / 255.),
        alf.layers.Conv2D(
            in_channels, 32, kernel_size=3, strides=2, padding=1),
        alf.layers.ResidueBlock(32, 32, 3, 1),
        alf.layers.ResidueBlock(32, 64, 3, 2),
        alf.layers.ResidueBlock(64, 64, 3, 1),
        nn.AvgPool2d(2),
        alf.layers.ResidueBlock(64, 64, 3, 1),
        nn.AvgPool2d(2),
        alf.layers.ResidueBlock(64, 64, 3, 1),
        alf.math.identity if use_bn else alf.math.normalize_min_max,
        input_tensor_spec=observation_spec,
    )


def create_dynamics_net(input_tensor_spec):
    state_spec, action_spec = input_tensor_spec
    plane_size = state_spec.shape[1:]
    num_planes = state_spec.shape[0]
    num_actions = action_spec.maximum + 1.
    return alf.nn.Sequential(
        lambda x: torch.cat([
            x[0], (x[1] / num_actions).reshape(-1, 1, 1, 1).expand(
                -1, 1, *plane_size)
        ],
            dim=1),
        a=alf.layers.Conv2D(
            num_planes + 1, 64, 3, padding=1, activation=alf.math.identity),
        b=(('input.0', 'a'), lambda x: (x[0] + x[1]).relu_()),
        c=alf.layers.ResidueBlock(64, 64, 3, 1),
        d=dyn_state_normalizer,
        input_tensor_spec=input_tensor_spec,
    )


def create_representation_net_small(observation_spec):
    size = observation_spec.shape[0]
    return alf.nn.Sequential(
        alf.layers.FC(size, 1024),
        input_tensor_spec=observation_spec,
    )


def create_dynamics_net_small(input_tensor_spec):
    state_spec, action_spec = input_tensor_spec
    plane_size = state_spec.shape[1:]
    num_planes = state_spec.shape[0]
    num_actions = action_spec.maximum + 1.
    return alf.nn.Sequential(
        lambda x: torch.cat([
            x[0], (x[1] / num_actions).reshape(-1, 1, 1, 1).expand(
                -1, 1, *plane_size)
        ],
            dim=1),
        a=alf.layers.Conv2D(
            num_planes + 1, 64, 3, padding=1, activation=alf.math.identity),
        b=(('input.0', 'a'), lambda x: (x[0] + x[1]).relu_()),
        c=dyn_state_normalizer,
        input_tensor_spec=input_tensor_spec,
    )


class SumNet(alf.nn.Network):
    def __init__(self, input_tensor_spec):
        super().__init__(input_tensor_spec, input_tensor_spec)

    def forward(self, inputs, state):
        state = inputs + state
        return state, state


def _get_rv_bias_initializer():
    if rv_bias_zero_init:
        return None
    return rv_loss.initialize_bias


@alf.configurable
def create_prediction_net(state_spec, action_spec, initial_game_over_bias=-5):
    dim = 32

    def _summarize_grad(x, name):
        if not x.requires_grad:
            return x
        if alf.summary.should_record_summaries():
            return summarize_tensor_gradients(
                "SimpleMCTSModel/" + name, x, clone=True)
        else:
            return x

    def _make_trunk(lstm=False):
        if lstm:
            if use_small_net:
                return [
                    alf.layers.Reshape(-1),
                    alf.nn.LSTMCell(7 * 7 * 64, 512),
                    alf.layers.FC(
                        512, dim, activation=torch.relu_, use_bn=pred_use_bn)
                ]
            else:
                return [
                    alf.layers.Conv2D(64, 16, 1),
                    alf.layers.Reshape(-1),
                    alf.nn.LSTMCell(6 * 6 * 16, 512),
                    alf.layers.FC(
                        512, dim, activation=torch.relu_, use_bn=pred_use_bn)
                ]
        else:
            if use_small_net:
                return [
                    alf.layers.Reshape(-1),
                    alf.layers.FC(
                        1,
                        1024,
                        activation=torch.relu_,
                        use_bn=use_bn),
                    alf.layers.FC(
                        1024, dim, activation=torch.relu_, use_bn=pred_use_bn),
                ]
            else:
                return [
                    alf.layers.ResidueBlock(64, 64, 3, 1),
                    alf.layers.Conv2D(64, 16, 3, padding=1),
                    alf.layers.Reshape(-1),
                    alf.layers.FC(
                        6 * 6 * 16,
                        dim,
                        activation=torch.relu_,
                        use_bn=pred_use_bn),
                ]

    if num_quantiles == 1:
        reshape_layer = [alf.layers.Reshape(())]
        reward_spec = alf.TensorSpec(())
    else:
        reshape_layer = []
        reward_spec = alf.TensorSpec((num_quantiles,))

    def _scale_grad(scale):
        if scale == 1.0:
            return []
        else:
            return [lambda x: torch.lerp(x.detach(), x, scale)]

    value_net = alf.layers.Sequential(
        partial(_summarize_grad, name='value_grad'),
        *_make_trunk(),
        *_scale_grad(1 / num_quantiles if scale_grad_by_num_quantiles else 1),
        alf.layers.FC(
            dim,
            num_quantiles,
            weight_opt_args=dict(weight_decay=rv_weight_decay),
            bias_opt_args=dict(weight_decay=rv_bias_decay),
            bias_initializer=_get_rv_bias_initializer(),
            kernel_initializer=torch.nn.init.zeros_),
        *_scale_grad(num_quantiles if scale_grad_by_num_quantiles else 1),
        *reshape_layer,
    )

    reward_net = [
                     partial(_summarize_grad, name='reward_grad'),
                     *_make_trunk(lstm_reward),
                     *_scale_grad(1 / num_quantiles if scale_grad_by_num_quantiles else 1),
                     alf.layers.FC(
                         dim,
                         num_quantiles,
                         weight_opt_args=dict(weight_decay=rv_weight_decay),
                         bias_opt_args=dict(weight_decay=rv_bias_decay),
                         bias_initializer=_get_rv_bias_initializer(),
                         kernel_initializer=torch.nn.init.zeros_),
                     *_scale_grad(num_quantiles if scale_grad_by_num_quantiles else 1),
                 ] + reshape_layer
    if sum_over_reward_prediction:
        reward_net.append(SumNet(reward_spec))
    reward_net = alf.nn.Sequential(*reward_net, input_tensor_spec=state_spec)

    action_net = alf.layers.Sequential(
        partial(_summarize_grad, name='policy_grad'),
        *_make_trunk(),
        alf.nn.CategoricalProjectionNetwork(
            dim,
            action_spec,
            logits_init_output_factor=0.0,
            weight_opt_args=dict(weight_decay=action_weight_decay),
            bias_opt_args=dict(weight_decay=action_bias_decay)),
    )
    game_over_net = alf.layers.Sequential(
        *_make_trunk(),
        alf.layers.FC(
            dim,
            1,
            kernel_initializer=torch.nn.init.zeros_,
            bias_init_value=initial_game_over_bias),
        alf.layers.Reshape(()),
    ) if train_game_over_function else lambda x: ()

    return alf.nn.Branch(
        value_net,
        reward_net,
        action_net,
        game_over_net,
        input_tensor_spec=state_spec,
    )


@alf.configurable
def repr_projection_net_ctor(input_tensor_spec,
                             hidden_size=512,
                             output_size=1024,
                             last_use_bn=False):
    return alf.nn.Sequential(
        alf.layers.Reshape(-1),
        alf.layers.FC(
            input_tensor_spec.numel,
            hidden_size,
            activation=torch.relu_,
            use_bn=True),
        alf.layers.FC(
            hidden_size, hidden_size, activation=torch.relu_, use_bn=True),
        alf.layers.FC(hidden_size, output_size, use_bn=last_use_bn),
        input_tensor_spec=input_tensor_spec)


@alf.configurable
def repr_prediction_net_ctor(input_tensor_spec,
                             hidden_size=512,
                             output_size=1024):
    return alf.nn.Sequential(
        alf.layers.FC(
            input_tensor_spec.numel,
            hidden_size,
            activation=torch.relu_,
            use_bn=True), alf.layers.FC(hidden_size, output_size))


encoding_net_ctor = create_representation_net_small
dynamics_net_ctor = create_dynamics_net_small

AdamOptimizers = {'ADAM': Adam, 'ADAMTF': AdamTF, 'ADAMW': AdamW}
lr_schedule = StepScheduler("percent", [(0.45, initial_lr),
                                        (0.9, 0.1 * initial_lr),
                                        (1.0, 0.01 * initial_lr)])

alf.config(
    "MCTSModel",
    value_loss=rv_loss,
    reward_loss=rv_loss,
    predict_reward_sum=True,
    policy_loss_weight=1.0,
    value_loss_weight=0.05,
    repr_prediction_loss_weight=40.0,
    reward_loss_weight=2.0)

alf.config(
    "SimpleMCTSModel",
    encoding_net_ctor=encoding_net_ctor,
    dynamics_net_ctor=dynamics_net_ctor,
    prediction_net_ctor=create_prediction_net,
    train_repr_prediction=train_repr_prediction,
    train_game_over_function=train_game_over_function,
    repr_projection_net_ctor=repr_projection_net_ctor,
    repr_prediction_net_ctor=repr_prediction_net_ctor,
    initial_alpha=0.)

alf.config(
    "MCTSAlgorithm",
    discount=discount,
    num_simulations=52,
    num_parallel_sims=2,
    root_dirichlet_alpha=0.3,
    root_exploration_fraction=0.,
    pb_c_init=0.75,
    pb_c_base=19652,
    is_two_player_game=False,
    visit_softmax_temperature_fn=VisitSoftmaxTemperatureByProgress(
        [(0.5, 1.0), (0.75, 0.5), (1, 0.25)]),
    act_with_exploration_policy=True,
    learn_with_exploration_policy=True,
    search_with_exploration_policy=True,
    unexpanded_value_score='mean',
    expand_all_children=False,
    expand_all_root_children=False,
    max_unroll_length=5,
    learn_policy_temperature=1.0)

alf.config(
    "MuzeroAlgorithm",
    enable_amp=True,
    mcts_algorithm_ctor=MCTSAlgorithm,
    model_ctor=SimpleMCTSModel,
    # The following line can be commented out if GPU memory is large enough
    reanalyze_batch_size=1280 if use_small_net else 640,
    num_unroll_steps=5,
    td_steps=10,
    reward_transformer=reward_transformer,
    reanalyze_td_steps_func=  # LinearMaxAgeTdStepFunc(),
    LinearTdStepFunc(max_bootstrap_age=1.2, min_td_steps=1),
    train_repr_prediction=train_repr_prediction,
    train_game_over_function=train_game_over_function,
    reanalyze_ratio=1.0,
    target_update_period=400,
    target_update_tau=1.0)

opt_kwargs = dict(
    lr=lr_schedule,
    weight_decay=weight_decay,
    gradient_clipping=1e9,
    clip_by_global_norm=True)

if optimizer_type in AdamOptimizers:
    optimizer = AdamOptimizers[optimizer_type](
        betas=(0.9, 0.999), eps=1e-7, **opt_kwargs)
else:
    optimizer = SGD(momentum=0.9, **opt_kwargs)

alf.config("Agent", optimizer=optimizer)

# training config
alf.config(
    "TrainerConfig",
    random_seed=2,
    unroll_length=unroll_length,
    mini_batch_size=256,
    num_updates_per_train_iter=5,
    update_counter_every_mini_batch=False,
    priority_replay=True,
    priority_replay_alpha=1.2,
    priority_replay_beta=LinearScheduler("percent", [(0.0, 0.4), (1.0, 1.0)]),
    num_iterations=0,
    num_env_steps=100000,
    num_checkpoints=1,
    evaluate=False,
    enable_amp=False,
    debug_summaries=True,
    summary_interval=int(100 // (unroll_length * num_envs)),
    replay_buffer_length=100000 // num_envs,
    initial_collect_steps=2000,
    summarize_grads_and_vars=True)

[emailweixu]

The three networks (encoding_net_ctor, dynamics_net_ctor, prediction_net_ctor) need to be changed for cartpole. The current version assumes image as input.

[ipsec]

Hi @emailweixu,

I think I'm in the right way.
I put CartPole-v1 to run with the config below.
My problem now is discovery the reason of it's not stable.

the config

import alf.examples.muzero_conf
import alf.layers
import torch
from alf.algorithms.mcts_algorithm import MCTSAlgorithm, VisitSoftmaxTemperatureByProgress
from alf.algorithms.mcts_models import SimpleMCTSModel
from alf.optimizers import AdamTF
from alf.utils import losses


def define_config(name, default_value):
    alf.define_config(name, default_value)
    return alf.get_config_value('_CONFIG._USER.' + name)


train_repr_prediction = define_config('train_repr_prediction', False)
train_game_over_function = define_config('train_game_over_function', False)
num_envs = define_config('num_envs', 5)
unroll_length = define_config('unroll_length', 10)
discount = define_config('discount', 0.99)
initial_lr = define_config('initial_lr', 5e-4)
terminal_on_life_loss = define_config('terminal_on_life_loss', False)
use_bn = define_config('use_bn', False)
pred_use_bn = define_config('pred_use_bn', False)
weight_decay = define_config('weight_decay', 0.1)
rv_weight_decay = define_config('rv_weight_decay', 4e-6)
rv_bias_decay = define_config('rv_bias_decay', 0)
action_weight_decay = define_config('action_weight_decay', None)
action_bias_decay = define_config('action_bias_decay', None)

rv_loss = define_config('rv_loss', losses.OrderedDiscreteRegressionLoss(
    transform=alf.math.Sqrt1pTransform(),
    inverse_after_mean=False))
rv_bias_zero_init = define_config('rv_bias_zero_init', False)

# This option can make the optimization of the parameters of quantile regression
# invariant to the number of quantiles.
scale_grad_by_num_quantiles = define_config('scale_grad_by_num_quantiles', False)

# alf.math.identity, alf.layers.BatchNorm2d(64), alf.math.normalize_min_max
dyn_state_normalizer = define_config('dyn_state_normalizer', alf.math.identity)

# If 1, not use quantile regression
num_quantiles = define_config('num_quantiles', 1)
alf.config('multi_quantile_huber_loss', delta=0.0)

alf.config(
    "create_environment",
    env_name="CartPole-v1",
    num_parallel_environments=num_envs)

support_size = define_config('support_size', 40)
hidden_size = define_config('hidden_size', 64)

def create_representation(observation_spec):
    return alf.nn.Sequential(
        alf.layers.FC(observation_spec.shape[0], hidden_size),
        alf.layers.FC(hidden_size, support_size),
        input_tensor_spec=observation_spec,
    )


@alf.configurable
def create_prediction_net(state_spec, action_spec, initial_game_over_bias=-5):

    value_net = alf.layers.Sequential(
        alf.layers.FC(state_spec.numel, hidden_size),
        alf.layers.FC(
            hidden_size,
            support_size,
            weight_opt_args=dict(weight_decay=rv_weight_decay),
            bias_opt_args=dict(weight_decay=rv_bias_decay),
            bias_initializer=rv_loss.initialize_bias,
            kernel_initializer=torch.nn.init.zeros_,
        ),
    )

    reward_net = alf.nn.Sequential(
        alf.layers.FC(state_spec.numel, hidden_size),
        alf.layers.FC(
            hidden_size,
            support_size,
            weight_opt_args=dict(weight_decay=rv_weight_decay),
            bias_opt_args=dict(weight_decay=rv_bias_decay),
            bias_initializer=rv_loss.initialize_bias,
            kernel_initializer=torch.nn.init.zeros_,
        ),
        input_tensor_spec=state_spec
    )

    action_net = alf.layers.Sequential(
        alf.layers.FC(state_spec.numel, hidden_size),
        alf.nn.CategoricalProjectionNetwork(
            hidden_size,
            action_spec,
            logits_init_output_factor=0.0,
            weight_opt_args=dict(weight_decay=action_weight_decay),
            bias_opt_args=dict(weight_decay=action_bias_decay)
        ),
    )

    game_over_net = alf.layers.Sequential(
        alf.layers.FC(state_spec.numel, hidden_size),
        alf.layers.FC(
            hidden_size,
            1,
            kernel_initializer=torch.nn.init.zeros_,
            bias_init_value=initial_game_over_bias,
        ),
        alf.layers.Reshape(()),
    )

    return alf.nn.Branch(
        value_net,
        reward_net,
        action_net,
        game_over_net,
        input_tensor_spec=state_spec,
    )


@alf.configurable
def repr_projection_net_ctor(input_tensor_spec,
                             # hidden_size=64,
                             output_size=support_size,
                             last_use_bn=use_bn):
    return alf.nn.Sequential(
        alf.layers.FC(
            input_tensor_spec.numel,
            hidden_size,
            activation=torch.relu_,
            use_bn=use_bn),
        alf.layers.FC(hidden_size, output_size, use_bn=last_use_bn),
        input_tensor_spec=input_tensor_spec)


@alf.configurable
def repr_prediction_net_ctor(input_tensor_spec,
                             # hidden_size=64,
                             output_size=support_size):
    return alf.nn.Sequential(
        alf.layers.FC(
            input_tensor_spec.numel,
            hidden_size,
            activation=torch.relu_,
            use_bn=use_bn),
        alf.layers.FC(hidden_size, output_size))


alf.config(
    "MCTSModel",
    value_loss=rv_loss,
    reward_loss=rv_loss,
    # predict_reward_sum=True,
    # policy_loss_weight=1.0,
    # value_loss_weight=0.05,
    # repr_prediction_loss_weight=40.0,
    # reward_loss_weight=2.0
)

alf.config(
    "SimpleMCTSModel",
    encoding_net_ctor=create_representation,
    prediction_net_ctor=create_prediction_net,
    train_repr_prediction=train_repr_prediction,
    train_game_over_function=train_game_over_function,
    repr_projection_net_ctor=repr_projection_net_ctor,
    repr_prediction_net_ctor=repr_prediction_net_ctor,
    # num_sampled_actions=20
    # initial_alpha=0.
)

alf.config(
    "MCTSAlgorithm",
    discount=discount,
    num_simulations=10,
    # num_parallel_sims=2,
    root_dirichlet_alpha=0.25,
    root_exploration_fraction=0.25,
    pb_c_init=1.25,
    pb_c_base=19652,
    is_two_player_game=False,
    visit_softmax_temperature_fn=VisitSoftmaxTemperatureByProgress(),
    act_with_exploration_policy=True,
    learn_with_exploration_policy=True,
    search_with_exploration_policy=True,
    expand_all_children=False,
    expand_all_root_children=True)

alf.config(
    "MuzeroAlgorithm",
    enable_amp=True,
    mcts_algorithm_ctor=MCTSAlgorithm,
    model_ctor=SimpleMCTSModel,
    num_unroll_steps=5,
    td_steps=10,
    train_repr_prediction=train_repr_prediction,
    train_game_over_function=train_game_over_function,
    reanalyze_ratio=1.0,
    target_update_period=1,
    target_update_tau=0.01)

alf.config("Agent", optimizer=AdamTF(lr=initial_lr))

# training config
alf.config(
    "TrainerConfig",
    unroll_length=unroll_length,
    mini_batch_size=256,
    num_updates_per_train_iter=10,
    # priority_replay=True,
    # priority_replay_alpha=1.2,
    # priority_replay_beta=LinearScheduler("percent", [(0.0, 0.4), (1.0, 1.0)]),
    num_iterations=10000,
    num_checkpoints=1000,
    evaluate=False,
    enable_amp=True,
    summary_interval=1,
    replay_buffer_length=10000,
    initial_collect_steps=100)

[emailweixu]

Some of the hyperparameters are tuned for Atari 100k setting (e.g. weight decay). And if you already have a working cartpole config, you can slightly change it use the new loss without introducing other changes (e.g., train_repr_prediction). Though train_repr_prediction is very useful for Atari, it's never tested on cartpole.

[ipsec]

Hi @emailweixu,

I'm running the debug of this version with categorical cross entropy loss function and I'm a bit confuse.

Using the muzero_pendulum_conf.py config (it use default loss function without categorical) the state passed to the _predict method on mcts_models.py has the shape torch.Size([1, 256]). This is right because of the batch_size=256.

While using the muzero_atari_conf.py version with categorical cross entropy loss the shape in the same point (_predict method) is torch.Size([2, 64, 6, 6]).

I thought it's could be the peculiarity of image shape (from atari), but using my custom_cartpole_config.py (bellow) I get the shape torch.Size([5, 30])

Where is the batch dimension of the state when using a custom model and different loss function?

Maybe my custom prediction, encoding or dynamics models could be wrong, but the same behavior occour with the atari version.

Where could I be failing?

Sorry by the newbie question.

Best regards,
Fernando

import torch
import alf
import alf.examples.muzero_conf
from alf.algorithms.data_transformer import RewardNormalizer
from alf.algorithms.mcts_algorithm import MCTSAlgorithm, VisitSoftmaxTemperatureByProgress
from alf.algorithms.mcts_models import SimpleMCTSModel, get_unique_num_actions
from alf.networks import EncodingNetwork, StableNormalProjectionNetwork, CategoricalProjectionNetwork
from alf.optimizers import Adam
from alf.utils.losses import OrderedDiscreteRegressionLoss


def define_config(name, default_value):
    alf.define_config(name, default_value)
    return alf.get_config_value('_CONFIG._USER.' + name)


train_repr_prediction = define_config('train_repr_prediction', False)
# initial_scale = define_config('initial_scale', 1.0)

alf.config(
    "create_environment",
    env_name="CartPole-v1",
    num_parallel_environments=5)

support_size = define_config('support_size', 30)
hidden_size = define_config('hidden_size', 64)


rv_loss = define_config(
    'rv_loss',
    OrderedDiscreteRegressionLoss(
        transform=alf.math.Sqrt1pTransform(), inverse_after_mean=False))
rv_bias_zero_init = define_config('rv_bias_zero_init', False)


@alf.configurable
def create_encoding_net(observation_spec):
    def _make_trunk():
        return [
            alf.layers.FC(
                observation_spec.numel,
                hidden_size,
                activation=torch.relu_),
            alf.layers.FC(
                hidden_size, hidden_size, activation=torch.relu_),
        ]

    encoding_net = [
        *_make_trunk(),
        alf.layers.FC(hidden_size, support_size, activation=torch.relu_, ),
    ]
    return alf.nn.Sequential(*encoding_net, alf.math.normalize_min_max)
    # return alf.nn.Sequential(*encoding_net)


@alf.configurable
def create_dynamics_net(input_tensor_spec):
    action_spec = input_tensor_spec[1]
    preproc = None
    if not action_spec.is_continuous:
        preproc = alf.layers.Sequential(
            alf.layers.OneHot(num_classes=get_unique_num_actions(action_spec)),
            alf.layers.Reshape([-1]))
    net = EncodingNetwork(
        input_tensor_spec,
        input_preprocessors=(None, preproc),
        preprocessing_combiner=alf.nest.utils.NestConcat(),
        fc_layer_params=(256, 256),
        # last_layer_size=input_tensor_spec[0].numel,
        last_layer_size=support_size,
        last_activation=torch.relu_
    )
    return alf.nn.Sequential(net, alf.math.normalize_min_max)


@alf.configurable
def create_prediction_net(state_spec, action_spec, initial_game_over_bias=0.0):
    def _make_trunk():
        return [
            alf.layers.FC(
                state_spec.numel,
                hidden_size,
                activation=torch.relu_),
            alf.layers.FC(
                hidden_size, hidden_size, activation=torch.relu_),
        ]

    value_net = [
        *_make_trunk(),
        alf.layers.FC(hidden_size, support_size, activation=torch.relu_, ),
    ]
    value_net = alf.layers.Sequential(*value_net)

    reward_net = [
        *_make_trunk(),
        alf.layers.FC(hidden_size, support_size, activation=torch.relu_, ),
    ]
    reward_net = alf.nn.Sequential(*reward_net, input_tensor_spec=state_spec)

    if action_spec.is_continuous:
        common_action_net = StableNormalProjectionNetwork(
            input_size=hidden_size,
            action_spec=action_spec
        )
    else:
        common_action_net = CategoricalProjectionNetwork(
            input_size=hidden_size,
            action_spec=action_spec
        )

    action_net = alf.layers.Sequential(
        *_make_trunk(),
        common_action_net,
    )

    game_over_net = alf.layers.Sequential(
        *_make_trunk(),
        alf.layers.FC(hidden_size, 1, activation=torch.relu_, ),
        alf.layers.Reshape(()),
    )

    return alf.nn.Branch(
        value_net,
        reward_net,
        action_net,
        game_over_net,
        input_tensor_spec=state_spec,
    )


alf.config(
    "MCTSModel",
    value_loss=rv_loss,
    reward_loss=rv_loss,
)

alf.config(
    "SimpleMCTSModel",
    train_repr_prediction=train_repr_prediction,
    encoding_net_ctor=create_encoding_net,
    dynamics_net_ctor=create_dynamics_net,
    prediction_net_ctor=create_prediction_net,
    # num_sampled_actions=20
)

alf.config(
    "MCTSAlgorithm",
    discount=0.99,
    num_simulations=10,
    root_dirichlet_alpha=0.25,
    root_exploration_fraction=0.25,
    pb_c_init=1.25,
    pb_c_base=19652,
    is_two_player_game=False,
    visit_softmax_temperature_fn=VisitSoftmaxTemperatureByProgress(),
    act_with_exploration_policy=True,
    learn_with_exploration_policy=True,
    search_with_exploration_policy=True,
    unexpanded_value_score='mean',
    expand_all_children=False,
    expand_all_root_children=True)

alf.config(
    "MuzeroAlgorithm",
    mcts_algorithm_ctor=MCTSAlgorithm,
    model_ctor=SimpleMCTSModel,
    num_unroll_steps=5,
    td_steps=10,
    reward_transformer=RewardNormalizer(update_mode="rollout"),
    train_repr_prediction=train_repr_prediction,
    reanalyze_ratio=1.0,
    #     target_update_period=1,
    #     target_update_tau=0.01
)

alf.config("Agent", optimizer=Adam(lr=1e-3))

# training config
alf.config(
    "TrainerConfig",
    unroll_length=10,
    mini_batch_size=256,
    num_updates_per_train_iter=10,
    num_iterations=100000,
    num_checkpoints=5,
    evaluate=False,
    summary_interval=0,
    num_summaries=100,
    replay_buffer_length=100000,
    initial_collect_steps=1000)

[emailweixu]

@ipsec The batch dimension is the first dimension. For rollout, it is same as num_parallel_environments, which is 1, 2, 5 in each of the three cases you mentioned.