[Feature] multiagent data standardization: PPO advantages

[matteobettini]

Hello!

Context

So, in this library we frequently apply standardizations over data (also called normalization) where we subtract the mean from the data and divide it by a standard deviation.

A few places where we do this that come to mind are in the transforms (where it can be used over observations and rewards) and in the advantages in ppo. Please let me know about other places if you can think of more.

Issue

The issue is that normally the mean and std statistics are computed over all the dimensions. But in multiagent settings, where for example a reward tensor could have an "agent" dimension, this is not the proper way to do it as different agents might have different reward distributions and applying standardization naively across them will create issues.

Example

For example consider the case where we have 2 agents each with 2 rewards (agent_1_r = [100,200], agent_2_r = [0,-1]), structured in a 2x2 tensor r = torch.tensor([[100,200],[0,-1]])

If we standardize using mean(r) and std(r), we get r = torch.tensor([[0.3,1.5],[-0.9,-.91]]), which is bad

If we standardize using mean(r, dim=1) and std(r,dim=1), hence excluding the agent dim from the statistics, we get r = torch.tensor([[-1,1],[1,-1]]), which is better as it does not cross pollute over agents

Proposal

We should adapt all places in the library where standardization is used to take a list of dimensions to exlude from the statistics. This way, multiagent users can provide the agent dimension as input and use the standardization feature

What I did here

I did this for the advantages in PPO

To do it, I made a nice standardization util function that takes as input a list of dimensions to exlude from the standardization (a more flexible implementation of a functional layer norm)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2677

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Failures because of deprecated version of actions/download-artifact: v3

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

I can definitely see why this is a problem, thanks for pointing it out and coming up with a solution!

My first thought when looking at this PR is that it can easily be forgotten.
IMO the ideal API is, by decreasing order of desiderability

1. The API doesn't break with edge cases (I'm using a overly generic term here, not saying MARL is edge case :) ) but takes care of them automatically. That can be achieved with a trainer class for instance.
2. The API breaks when encoutering an edge case and the error message tells the user exactly why and what to do
3. The run goes through but a warning is raised.
   Finally, 4., the run goes through and things are silently wrong.

I think this is a (4). The ideal solution would be to do it automatically (1) but I'm not sure how we can achieve that with the current API. Any thoughts on how to make it a (2) or (3)?
(As a counter-example, if your data has some specific keys, the loss will simply not work until you call loss_module.set_keys(**tensor_keys). We're not doing it automatically but there is very little chance that this step is forgotten and we can capture the errors and give an informative error message to let people know what to do.)

Another way of saying this is: there are two levels of complexity that make a lib hard to use and I think both are present with the current solution:

1. (weak) There is something that users need to know when they use the library but they may not be aware of that and the lib isn't telling them at runtime;
2. (strong) The lib doesn't take care automaticlly of edge cases and it's up to the user to come up with a solution.
   How can we make sure that users are not assuming that this is done automatically when they say "normalize_advantage=True`? Is there a quick check that we can do?

One thing I have in mind is to look at the reward shape and detecting if the last dim is a singleton (if not, the reward is probably something that isn't from regular single agent)? In principle, anything with shape that is not (*tensordict.shape, 1) should have a dedicated reward normalization, no?

Regarding other parts of the lib where this could be needed:

• I think we also need to address this issue with the advantages where there's a lot of normalization of rewards etc.
  

  rl/torchrl/objectives/value/advantages.py

  Lines 1434 to 1438 in d009835

  	if self.average_gae:
  	loc = adv.mean()
  	scale = adv.std().clamp_min(1e-4)
  	adv = adv - loc
  	adv = adv / scale

  
  

  rl/torchrl/objectives/value/advantages.py

  Lines 1804 to 1808 in d009835

  	if self.average_adv:
  	loc = adv.mean()
  	scale = adv.std().clamp_min(1e-5)
  	adv = adv - loc
  	adv = adv / scale

• VecNorm needs a big refactoring (it's clunky and inefficient atm), so we can put this in the TODO list.
• For ObservationNorm, you could pass a loc / scale with any broadcastable shape and it should work.
  The ObservationNorm.init_stats should also have the required keyword arguments to account for the example you gave.

[vmoens]

Suggested change

	normalize_advantage_exclude_dims (Sequence[int], optional): dimensions to exclude from the advantage
	normalize_advantage_exclude_dims (Tuple[int], optional): dimensions to exclude from the advantage

[vmoens]

This seems like an indirect argument to me and it could be hard for users to understand.
If you're using PPO with this kwarg, you need to think this way:

1. I asked to normalize the advantage
2. My reward has one or more agent dims
3. I want to excldue these from the advantage normalization.
   Another way to do this could be to have a kwarg

reward_agent_dim: int | Tuple[int] | None = None,

If this is passed, the reward_agent_dim will be kept during calls to mean / std.
Per se, it changes nothing but it's a tiny bit more clear from the name and signature that this indicates a dim that is not a feature dim.

[matteobettini]

See my discussion point on this being useful for more uscases. We can come up with a better name together

[vmoens]

Suggested change

	standardization, can be negative. Useful in multiagent settings to exlude the agent dimension. Default: ().
	standardization. Negative dimensions are valid. This is useful in multiagent settings where the agent dimension may be excluded from the reductions.
	Defaults to ``()``.

[vmoens]

I have a question reguarding the "negative" dims.

I guess ultimately we would also like to do the same with this

rl/torchrl/objectives/value/advantages.py

Lines 696 to 701 in d009835

	if self.average_rewards:
	reward = reward - reward.mean()
	reward = reward / reward.std().clamp_min(1e-5)
	tensordict.set(
	("next", self.tensor_keys.reward), reward
	) # we must update the rewards if they are used later in the code

But there we have an extra "time" dimension, so we must be careful that agent_dim=(-2,) may actually mean (-3,) in some cases (similarly, feature_dim=(4,) may be feature_dim=(5,)).

Is it an issue in this (ie, PPO) case?

[matteobettini]

I think the user should pass the dimension index for the tensor that goes in the class.

If you give as input [batch,time,agents,n] then agents is -2

also agents is usually after time

[matteobettini]

Thanks so much for the in-depth reveiw! I agree!

Let me trying t give some answers, I'll start from some general clarifications

This is not a feature purely for multiagent uses. The need for this feature could be various. Imagine a multi-objective rl setting where you have an array of rewards. Now, even in that case you would not want to mix rewards during normalization. You could also have mutli-objective + multi-agent. So, the goal of this PR is more like saying that using tensor.mean() in the lib is a tight fit as there are a lot of usecases where users might want to keep some dimensions independent while computing statistics. So the idea is to provide a way for users to communicate that, regardelss of the usecase. For this, I don't think that naming it something like "agent_dimension" is appropriate as the poor multi-objective usecases or others are not represented

I agree this should be automatic. The problem is that it is quite impossible to automatically know what dimensions the user wants to keep independent. What we could do to turn it into a (3) is what you suggest: if you are asking for normalization and your shape is more than (*tensordict.shape, 1), we can give you a warning saying "hey! you have all those extra dims that we are considering in our statistics, any of those independent?"

[matteobettini]

I should have addressed all comments.

Anorther thing I forgot to say is that the _standardize ultil function is meant to be used in all those other cases we need to fix you mentioned