Something isn't right with actions (and chance) encoding

[Firerozes]

Hello,

Something isn't right with actions/chance encoding. I only checked for muzero and stochastic muzero (lzero/model), but these 4 follow the same overall process - and I guess other algos too. This one is taken from lzero/model/muzero_model_mlp :

        # discrete action space
        if self.discrete_action_encoding_type == 'one_hot':
            # Stack latent_state with the one hot encoded action
            if len(action.shape) == 1:
                # (batch_size, ) -> (batch_size, 1)
                # e.g.,  torch.Size([8]) ->  torch.Size([8, 1])
                action = action.unsqueeze(-1)

            # transform action to one-hot encoding.
            # action_one_hot shape: (batch_size, action_space_size), e.g., (8, 4)
            action_one_hot = torch.zeros(action.shape[0], self.action_space_size, device=action.device)
            # transform action to torch.int64
            action = action.long()
            action_one_hot.scatter_(1, action, 1)
            action_encoding = action_one_hot
        elif self.discrete_action_encoding_type == 'not_one_hot':
            action_encoding = action / self.action_space_size
            if len(action_encoding.shape) == 1:
                # (batch_size, ) -> (batch_size, 1)
                # e.g.,  torch.Size([8]) ->  torch.Size([8, 1])
                action_encoding = action_encoding.unsqueeze(-1)

        action_encoding = action_encoding.to(latent_state.device).float()
        # state_action_encoding shape: (batch_size, latent_state[1] + action_dim]) or
        # (batch_size, latent_state[1] + action_space_size]) depending on the discrete_action_encoding_type.
        state_action_encoding = torch.cat((latent_state, action_encoding), dim=1)

        next_latent_state, reward = self.dynamics_network(state_action_encoding)

        if not self.state_norm:
            return next_latent_state, reward
        else:
            next_latent_state_normalized = renormalize(next_latent_state)
            return next_latent_state_normalized, reward

• If self.discrete_action_encoding_type == 'one_hot', there is some encoding...but it doesn't follow DeepMind's Zero algo's ideas. That is, there are several 1s that are spread in either the first or the second colum of a tensor. This doesn't match e.g. AlphaZero's philosophy of putting a single 1 in played position for Go
• If self.discrete_action_encoding_type != 'one_hot', there is some kind of subtitute, that is, action / self.action_space_size. This is, I guess, okayish for action space sizes of low cardinality. But for high ones ? E.g., for Go's 362 state, the difference between two states is ~0.003. Isn't that too small to have the model truly distinct actions ?

[puyuan1996]

Hello, this is an excellent and insightful question that highlights two critical aspects of action encoding in reinforcement learning models. Your observations are very sharp. Below, I will provide a detailed explanation for both of your concerns.

1. Concern about the 'one_hot' Encoding

You mentioned that the one-hot encoding implementation in the code doesn't align with AlphaZero's philosophy and that you observed "several 1s that are spread in either the first or the second column of a tensor." There might be a slight misunderstanding here; this code is, in fact, implementing a standard and correct one-hot encoding.

Let's break down the code:

# 1. Create a tensor of zeros with shape (batch_size, action_space_size)
action_one_hot = torch.zeros(action.shape[0], self.action_space_size, device=action.device)

# 2. Ensure the action tensor is of type long, as required by the scatter_ function
action = action.long() 

# 3. The core operation: scatter_
#    - dim=1: Indicates that the operation is performed along the second dimension (columns).
#    - index=action: This is a tensor of shape (batch_size, 1) that specifies the index (column) to place a value in for each row.
#    - src=1: This specifies that the value to be placed is 1.
action_one_hot.scatter_(1, action, 1)

The scatter_ function works as follows: for each sample in the batch (i.e., for each row), it places the value 1 at the column index specified by action. As a result, each row in the final action_one_hot tensor will have exactly one 1, with all other elements being 0.

Why might you have seen "several 1s"?
A likely reason is that when you printed or inspected the entire action_one_hot tensor, you were looking at a full batch of data. For instance, with a batch size of 8 and an action space of 4, you might see an output like this:

tensor([[0., 1., 0., 0.],  # Sample 1, action was 1
        [1., 0., 0., 0.],  # Sample 2, action was 0
        [0., 0., 0., 1.],  # Sample 3, action was 3
        [0., 1., 0., 0.],  # Sample 4, action was 1
        ...])

In this batch, the first column (index 0) and the second column (index 1) do indeed contain multiple 1s. However, the crucial point is that each row contains only a single 1. This perfectly aligns with the definition of one-hot encoding and is consistent with the philosophy of representing an action for a model's input (like in AlphaZero) in a clear, unambiguous way.

Therefore, this part of the implementation is correct and standard.

2. Concern about the 'not_one_hot' Encoding

Your concern about this encoding method is absolutely correct and very insightful.

Encoding the action as action / self.action_space_size is a form of normalized scalar encoding. Your analysis is spot-on:

• Low Numerical Separability: As you calculated, for an environment like Go with 362 possible actions, the encoded values for adjacent actions (e.g., action 10 and action 11) would be 10/362 ≈ 0.0276 and 11/362 ≈ 0.0304, respectively. The difference between them is tiny. This requires the initial layers of the neural network to be extremely sensitive to small numerical differences to distinguish their effects, which is very difficult in practice. The model is likely to treat these adjacent actions as "similar," failing to learn their potentially vast differences in outcome.

• Introduction of a False Ordinal Relationship: This encoding method imposes a linear, ordered relationship between actions. It tells the model that action 2 is "between" action 1 and action 3, and that the "distance" between action 100 and 101 is the same as the distance between action 1 and 2. For categorical actions like Go moves, this ordinal relationship does not exist (point A1 on the board is not numerically "less than" point A2). This introduces an unnecessary and incorrect inductive bias into the model's learning process.

Why does this option exist?
This 'not_one_hot' option might be acceptable in very specific scenarios or serve as a simplified baseline for research:

1. Extremely Small Action Spaces: For example, if there are only 2 or 3 actions.
2. Naturally Ordinal Actions: For instance, if the actions were "set speed to gear 1 / gear 2 / gear 3 / gear 4." In such a case, a numerical encoding could be meaningful.

In summary, your judgment is correct. For most discrete, categorical action spaces (especially large ones like Go), 'one_hot' encoding is the more robust, standard, and less problematic choice. The 'not_one_hot' method is a simplified approach with significant drawbacks and is not recommended for serious applications.

(Modified with assistance from Gemini-2.5-Pro.)

[Firerozes]

Hello,

I indeed made a careless mistake about 'one_hot' encoding. Even though this implementation can bring some differences with original DeepMind's ideas, the principal is there.

I also understand your points on 'not_one_hot' encoding.

As a conclusion, I have this question/suggestion :
Right now, you have in \lzero\model more than 10 models. All of them have common points. For example, all the MuZero algorithms encode the action and join this encoding with the latent state in a similar manner.
BUT it's rewritten in each verion. And then some contributors modify one of them. And not the others. And this brings chaos. And even right now, for example, stochastic_muzero_model.py only supports 'not_one_hot' and I'm almost certain it's just because updating it was forgotten !
Wouldn't it be better to solve this by having, for example, one (or a few) clean dynamics functions defined in \lzero\model\commons.py, and then defining dynamics in xxx_model.py as such :

from .common import dynamics_muzero     #or e.g; dynamics_lightzero
"""BLAH BLAH"""
     def __init__(
     """BLAH BLAH"""
          self._dynamics = dynamics_muzero

If you find such idea a good one, I could possibly do it on a branch. This is, according to me, standardization effort that would afterwards rationalize how updates are brought.

But I'd rather have your confirmation before starting such work and risking seeing my PRs not accepted.