Is there a way for tensordict_keys to be defined at the class level (and not instance)? Such that one can do PPOLoss.tensordict_keys without instantiating an object?

That would lead that the behavior of one instance of a loss module would change if .set_keys() is called somewhere else. I think that could surprise the user of torchrl. (Although, usually only one loss module is used.)

Also, what about making tensordict_keys a property using @abc.abstractmethod in LossModule parent class, such that we force all new losses to have that attribute?

That is a good idea, I will try to implement this. It should lead to an easier interface.

That would lead that the behavior of one instance of a loss module would change if .set_keys() is called somewhere else. I think that could surprise the user of torchrl. (Although, usually only one loss module is used.)

Not necessarily.
Here you clearly separate the default and non-default. The default should be defined at the class level if they're kept separated.

Not necessarily.
Here you clearly separate the default and non-default. The default should be defined at the class level if they're kept separated.

Got it. I've implemented this idea and it looks good. Although, I had to use a abstract staticmethod. Afaik, there is no static property decorator.

Also one thing to mention is the the new keys have to be transparently reflected in the value_estimators

loss = loss()
loss.set_keys()
loss.make_value_estim()

and

loss = loss()
loss.make_value_estim()
loss.set_keys()

Should both work with the keys reflected in the value estimators

Also one thing to mention is the the new keys have to be transparently reflected in the value_estimators

loss = loss()
loss.set_keys()
loss.make_value_estim()

and

loss = loss()
loss.make_value_estim()
loss.set_keys()

Should both work with the keys reflected in the value estimators

ideally keys of the value estimator should only be set through the value estimator no?

But the value estimator is created by the loss.

do we want to make users to call set_keys twice?

already calling it once is a added complexity for MARL users

But the value estimator is created by the loss.

do we want to make users to call set_keys twice?

already calling it once is a added complexity for MARL users

That is a good point. This problem is introduced by allowing to configure the tensordict keys after constructing the losses.
I could solve this by extending .set_keys(...) so that it also sets the keys for the value estimator.
Another solution would be to remove .set_keys(...) and add a generic constructor argument for all loss modules, e.g., tensordict_keys.
If there is no need to configure the keys during runtime, I would prefer the later one, although the code becomes a bit clunky.

The solution with the ctor would roughly look like:

class MyLoss(LossModule):
    def __init__(self, ...., tensordict_keys = {}):
         super().__init__(tensordict_keys=tensordict_keys)
         ...

class LossModule(nn.Module, ABC):
    def __init__(self, tensordict_keys={}):
        #merge default tensordict keys with provided keys
       default_keys = self.default_tensordict_keys()
       for key, value in tensordict_keys.items():
          if key not in default_keys:
              raise ValueError(...)
          default_keys[key] = value
      // create attributes
      ...

Forwarding the new keys via .set_keys() to the value estimator would also require to maintain a mapping from key_name_loss to key_name_value_estimator. In the default case both key names are identical, see ppo.py

# from ppo.py
def make_value_estimator(...):
    ...
    value_key = self.value_key
    if value_type == ValueEstimators.TD1:
        self._value_estimator = TD1Estimator(
            value_network=self.critic, value_key=value_key, **hp
        )

However, there is also this case ddpg.py:

# from ddpg.py
def make_value_estimator(...):
    ...
    value_key = "state_action_value" # <- would correspond to self.state_action_value_key
    if value_type == ValueEstimators.TD1:
        self._value_estimator = TD1Estimator(
            value_network=self.actor_critic, value_key=value_key, **hp
        )

Btw, this code will raise an exception in case the value key name has a non-default value

loss = loss()
loss.make_value_estim()
loss.set_keys()

because

actor = # construct tensordict with non-default value key
loss = loss() # constructing loss with default key names works because the ctor doesn't use any key names
loss.make_value_estim() # will raise an exception because

The value estimator checks that the value key is in the value_network.out_keys.
Since the actor uses a non-default value key, but loss and hence the value estimator are instantiated with the default values this will raise an exception.