AlphaZero not working on custom env: solving algebraic equations

[Khev]

Hi there,

CONTEXT

Many thanks for this great repo. I'm interested in adding an "equation solving" env, in which an RL agent must learn to solve equations like

ax + b = 0
a/x + b = 0
(ax+b)/(cx+d) + e = 0

By taking actions like (operation, term), where operation \in (add, sub, mul, div) and term \in (a,b,ax, ...) = terms in the equation. So sovling the equation ax +b = 0 requires two actions (sub, b), (div, a). I have solved most equations using ppo in about 10^3 or 10^4 steps.

ISSUE

I am trying to use your impementation of AlphaZero to solve the envs. But its struggling to solve even basic equation like

x +b = 0
ax+b = 0

The learning curves are below. You can see its getting the right solution pretty often, but its failing to learn from these experiences. The policy fails to converge. I'm wondering if there's something wrong with my setup. Below are my config files

Any insight would be most welcome.

One other question: would you be interested in adding this custom env to the repo? Its a challenging benchmark for RL based symbolic mathematics. It would be a good contribution for the community.

FILES

from easydict import EasyDict

# ==============================================================
# Frequently changed config specified by the user (lightweight settings)
# ==============================================================
collector_env_num = 4         # Number of parallel environments for data collection
n_episode = 4                 # Number of episodes per training iteration
evaluator_env_num = 1         # Number of evaluator environments
num_simulations = 50          # MCTS simulations per move (try increasing if needed)
update_per_collect = 100      # Number of gradient updates per data collection cycle
batch_size = 32               # Mini-batch size for training
max_env_step = int(1e5)       # Maximum total environment steps for a quick run
model_path = None
mcts_ctree = False

# ==============================================================
# Configurations for singleEqn_env (lightweight version)
# ==============================================================
singleEqn_alphazero_config = dict(
    exp_name='data_alphazero/singleEqn/a/x+b/',
    env=dict(
        battle_mode='play_with_bot_mode',
        battle_mode_in_simulation_env='self_play_mode',  # For simulation during MCTS
        channel_last=False,
        collector_env_num=collector_env_num,
        evaluator_env_num=evaluator_env_num,
        n_evaluator_episode=evaluator_env_num,
        manager=dict(shared_memory=False),
        agent_vs_human=False,
        prob_random_agent=0,
        prob_expert_agent=0,
        prob_random_action_in_bot=0,
        scale=True,
        render_mode=None,
        replay_path=None,
        alphazero_mcts_ctree=mcts_ctree,
    ),
    policy=dict(
        mcts_ctree=mcts_ctree,
        simulation_env_id='singleEqn_env',  # Must match the registered name of your environment
        model=dict(
            type='AlphaZeroMLPModel',
            import_names=['zoo.custom_envs.equation_solver.my_alphazero_mlp_model'],
            observation_shape=(41,),        # Flat vector of length 41
            action_space_size=50,             # Matches your environment's action_dim
            hidden_sizes=[64, 64],          # MLP hidden layer sizes
        ),
        cuda=True,
        env_type='not_board_games',
        action_type='varied_action_space',
        update_per_collect=update_per_collect,
        batch_size=batch_size,
        optim_type='Adam',
        lr_piecewise_constant_decay=False,
        # learning_rate=0.003,
        learning_rate=3e-4,
        grad_clip_value=0.5,
        value_weight=1.0,
        entropy_weight=0.0,
        n_episode=n_episode,
        eval_freq=int(2e3),
        mcts=dict(num_simulations=num_simulations),
        collector_env_num=collector_env_num,
        evaluator_env_num=evaluator_env_num,
        other=dict(
            replay_buffer=dict(
                type='advanced',              # Use advanced (or prioritized) replay buffer
                replay_buffer_size=10000,       # Set a smaller buffer for lightweight runs
                sample_min_limit_ratio=0.25,      # Allow sampling even if only 50% of batch size is available.
                alpha=0.6,
                beta=0.4,
                anneal_step=100000,
                enable_track_used_data=False,
                deepcopy=False,
                save_episode=False,
            )
        ),
    ),
)
singleEqn_alphazero_config = EasyDict(singleEqn_alphazero_config)
main_config = singleEqn_alphazero_config

singleEqn_alphazero_create_config = dict(
    env=dict(
        type='singleEqn_env',
        import_names=['zoo.custom_envs.equation_solver.env_single_eqn'],  # Adjust this path if needed
    ),
    env_manager=dict(type='subprocess'),
    policy=dict(
        type='MyAlphaZeroPolicy',  # Your custom policy subclass
        import_names=['zoo.custom_envs.equation_solver.my_alphazero_policy'],
    ),
    collector=dict(
        type='episode_alphazero',
        import_names=['lzero.worker.alphazero_collector'],
    ),
    evaluator=dict(
        type='alphazero',
        import_names=['lzero.worker.alphazero_evaluator'],
    )
)
singleEqn_alphazero_create_config = EasyDict(singleEqn_alphazero_create_config)
create_config = singleEqn_alphazero_create_config

if __name__ == '__main__':
    from lzero.entry import train_alphazero
    # Merge the environment configuration into the policy config.
    main_config.policy.env = main_config.env
    train_alphazero([main_config, create_config], seed=0, model_path=model_path, max_env_step=max_env_step)

[Khev]

More info: I also created a custom alphaZero model which uses MLP instead of a CNN, and a custom AlphaZero policy which defines my simulation envs. See below.

"""
A custom MLP-based AlphaZero model for vector inputs.
This model is designed for tasks where the state is a simple vector (e.g., [3, -2, 1, 10, ...]).
We support two observation shape formats:
  - A flat vector, e.g. (41,)
  - A pseudo-image, e.g. (1, 41, 1)
In either case, the input is flattened and passed through an MLP to produce both policy logits and a value.
"""

from typing import Tuple, Optional, Sequence
import torch
import torch.nn as nn
import torch.nn.functional as F
from ding.torch_utils import MLP  # DI-engine's MLP utility
from ding.utils import MODEL_REGISTRY, SequenceType

@MODEL_REGISTRY.register('AlphaZeroMLPModel')
class AlphaZeroMLPModel(nn.Module):
    def __init__(
        self,
        observation_shape: SequenceType = (41,), 
        action_space_size: int = 50,
        categorical_distribution: bool = False,
        activation: Optional[nn.Module] = nn.ReLU(inplace=True),
        hidden_sizes: SequenceType = [128, 128],
        last_linear_layer_init_zero: bool = True,
    ):
        """
        Args:
            observation_shape: Expected shape of observations.
                Can be a flat vector, e.g. (41,), or a pseudo-image, e.g. (1, 41, 1).
            action_space_size: Number of discrete actions.
            categorical_distribution: If True, use a categorical representation for value.
            activation: Activation function.
            hidden_sizes: List of hidden layer sizes for the shared MLP.
            last_linear_layer_init_zero: Whether to initialize the last layer of the heads with zeros.
        """
        super(AlphaZeroMLPModel, self).__init__()
        self.categorical_distribution = categorical_distribution
        self.observation_shape = observation_shape
        self.value_support_size = 601 if self.categorical_distribution else 1
        self.last_linear_layer_init_zero = last_linear_layer_init_zero

        # Determine the input dimension based on the observation shape.
        if len(observation_shape) == 1:
            self.input_dim = observation_shape[0]
        elif len(observation_shape) == 3:
            self.input_dim = observation_shape[0] * observation_shape[1] * observation_shape[2]
        else:
            raise ValueError(f"Unsupported observation_shape: {observation_shape}")

        self.action_space_size = action_space_size

        # Shared representation network: a simple MLP.
        self.representation_network = MLP(
            in_channels=self.input_dim,
            hidden_channels=hidden_sizes[0],
            out_channels=hidden_sizes[-1],
            layer_num=len(hidden_sizes),
            activation=activation,
        )

        # Policy head: maps shared representation to action logits.
        self.policy_head = MLP(
            in_channels=hidden_sizes[-1],
            hidden_channels=hidden_sizes[0],
            out_channels=action_space_size,
            layer_num=len(hidden_sizes),
            activation=activation,
            last_linear_layer_init_zero=last_linear_layer_init_zero,
        )

        # Value head: maps shared representation to a scalar value.
        self.value_head = MLP(
            in_channels=hidden_sizes[-1],
            hidden_channels=hidden_sizes[0],
            out_channels=self.value_support_size,
            layer_num=len(hidden_sizes),
            activation=activation,
            last_linear_layer_init_zero=last_linear_layer_init_zero,
        )

    def forward(self, state_batch: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            state_batch: Tensor of shape (B, C, H, W) or (B, L) if flat.
        Returns:
            Tuple (policy_logits, value) where:
              - policy_logits: (B, action_space_size)
              - value: (B, 1) (if not using categorical distribution)
        """
        B = state_batch.size(0)
        # If the state is 4D, flatten; if it's 2D, it's already flat.
        if state_batch.dim() == 4:
            x = state_batch.view(B, -1)
        else:
            x = state_batch
        rep = self.representation_network(x)
        policy_logits = self.policy_head(rep)
        value = self.value_head(rep)
        if not self.categorical_distribution:
            value = value.unsqueeze(-1)
        return policy_logits, value

    def compute_policy_value(self, state_batch: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Computes the softmax probabilities and state value.
        """
        logits, value = self.forward(state_batch)
        prob = F.softmax(logits, dim=-1)
        return prob, value

    def compute_logp_value(self, state_batch: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Computes the log-softmax probabilities and state value.
        """
        logits, value = self.forward(state_batch)
        log_prob = F.log_softmax(logits, dim=-1)
        return log_prob, value

from lzero.policy.alphazero import AlphaZeroPolicy
from ding.policy.base_policy import POLICY_REGISTRY
from ding.utils import ENV_REGISTRY
from easydict import EasyDict

@POLICY_REGISTRY.register('MyAlphaZeroPolicy')
class MyAlphaZeroPolicy(AlphaZeroPolicy):
    def _get_simulation_env(self):
        # Create a simulation environment using the merged environment config.
        sim_cfg = EasyDict(self.cfg.get('env', {}))
        sim_env = ENV_REGISTRY.get(self.cfg.simulation_env_id)(cfg=sim_cfg)
        return sim_env

    @property
    def simulate_env(self):
        # Lazily initialize and cache the simulation environment.
        if not hasattr(self, '_simulate_env'):
            self._simulate_env = self._get_simulation_env()
        return self._simulate_env

[puyuan1996]

We are highly interested in incorporating this symbolic mathematics environment into our repository. If you have the time, we would greatly appreciate it if you could submit a pull request (PR) following the format of our existing environments. Please ensure that the submission includes comprehensive documentation detailing the formal definition of the Markov Decision Process (MDP) components and their explanations.

Regarding the issue of AlphaZero failing to converge in this environment, we encourage you to submit a PR with the relevant implementation so that we can review the code in detail. Additionally, I will conduct local testing in the coming days to further investigate the issue. Please note that AlphaZero requires the environment to support reset(), which your environment appears to satisfy. However, if the environment does not support reset(), as is the case for certain single-player tasks, MuZero may be a more appropriate choice.

Thank you for your contribution, and we look forward to your PR.

[puyuan1996]

a custom AlphaZero policy which defines my simulation envs. See below.

What does 'a custom AlphaZero policy that defines my simulation environments' mean? My understanding is that the environment should be clearly defined within the environment module.

[Khev]

Hi there,

That's great! I will set up the PR shortly.