Defining a ProbalisticActor with two normal distributions

Mox_Antia: <p>Hi All,</p>
<p>Just starting using pyTorch and TorchRL, I am trying to adapt the <a href="https://pytorch.org/rl/stable/tutorials/coding_ppo.html" rel="noopener nofollow ugc">PPO tutorial</a> to my own settings.</p>
<p>I want to train robots which have a mixture of wheels and legs. So, I have a composite actor spec with joints and wheels output.<br>
But I fail to define a working probabilistic actor.</p>
<p>My actor_spec:</p>
<pre><code class="lang-auto">        self.action_spec = CompositeSpec(
            joints=BoundedTensorSpec(
                low=-torch.pi,
                high=torch.pi,
                shape=(len(self.robot.joint_ids),),
                dtype=torch.float32,
                device=self.device,
                ),
            wheels=BoundedTensorSpec(
                low=-100,
                high=100,
                shape=(len(self.robot.wheel_ids),),
                dtype=torch.float32,
                device=self.device,
           ) 
        )
</code></pre>
<p>And the definition of my actor:</p>
<pre><code class="lang-auto">actor_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(2 * env.action_spec["joints"].shape[-1] + 2 * env.action_spec["wheels"].shape[-1], device=device),
    NormalParamExtractor(),
)

# Define the policy module
policy_module = TensorDictModule(
    actor_net, in_keys=["observation"], 
    out_keys=["joints_loc", "joints_scale", "wheels_loc", "wheels_scale"]
)


actor = ProbabilisticActor(
    module=policy_module,
    spec=env.action_spec,
    in_keys=["joints_loc", "joints_scale", "wheels_loc", "wheels_scale"],
    out_keys=["joints", "wheels"],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "min": torch.cat([env.action_spec["joints"].space.low, env.action_spec["wheels"].space.low]),
        "max": torch.cat([env.action_spec["joints"].space.high, env.action_spec["wheels"].space.high]),
    },
)
</code></pre>
<p>When I run <code>actor(env.reset())</code> I get this error:</p>
<pre><code class="lang-auto">KeyError: 'key "wheels_loc" not found in TensorDict with keys [\'done\', \'joints_loc\', \'joints_scale\', \'observation\', \'step_count\', \'terminated\']'
</code></pre>
<p>It seems that it does only take into account the two first out_keys of my module.</p>
<p>How would you implement a probabilistic actor with a composite action output? Or how to define a probabilistic actor with two normal distribution?</p>
<p>Many thanks!</p>

vmoens: <p>Use a <code>CompositeDistribution</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">from torch import distributions as d
from tensordict import TensorDict
import torch
from tensordict.nn import CompositeDistribution

params = TensorDict({
     "joint": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
     "wheel": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
 }, [3])
dist = CompositeDistribution(params,
    distribution_map={"joint": d.Normal, "wheel": d.Normal})
sample = dist.sample((4,))
print(sample)
</code></pre>
<p>LMK if you need further help!</p>

Mox_Antia: <p>Thank you for the quick answer.<br>
So I integrated a CompositeDistribution following your suggestion:</p>
<pre><code class="lang-auto">actor_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(2 * env.action_spec["joints"].shape[-1] + 2 * env.action_spec["wheels"].shape[-1], device=device),
    NormalParamExtractor(),
)

params = TensorDict({
    "joints": {"loc": torch.rand(env.action_spec["joints"].shape[-1]), "scale": torch.rand(env.action_spec["wheels"].shape[-1])},
    "wheels": {"loc": torch.rand(env.action_spec["joints"].shape[-1]), "scale": torch.rand(env.action_spec["wheels"].shape[-1])}},
    batch_size=[],
)


# Define the policy module
policy_module = TensorDictModule(
    actor_net, in_keys=["observation"], 
    out_keys=[("joints","loc"), ("joints"),("scale"), ("wheels","loc"), ("wheels","scale")]
)


actor = ProbabilisticActor(
    module=policy_module,
    spec=env.action_spec,
    in_keys=[("joints","loc"), ("joints"),("scale"), ("wheels","loc"), ("wheels","scale")],
    out_keys=["joints", "wheels"],

    distribution_class=CompositeDistribution,
    distribution_kwargs={
        "params" : params,
        "distribution_map": {"joints": d.Normal, "wheels": d.Normal},
        "min": torch.cat([env.action_spec["joints"].space.low, env.action_spec["wheels"].space.low]),
        "max": torch.cat([env.action_spec["joints"].space.high, env.action_spec["wheels"].space.high]),
    },
)
</code></pre>
<p>Now I have this error:</p>
<pre><code class="lang-auto">AttributeError: 'BoundedTensorSpec' object has no attribute 'items'
</code></pre>
<p>I guess it comes from the action_spec and how I write my out_keys for the policy_module and in_keys for the ProbabilisticActor.</p>
<p>I am not sure how to write them the right way.</p>

vmoens: <p>Here is the corrected code</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torchrl.data import ReplayBuffer, LazyMemmapStorage
from torch import distributions as d
from tensordict.nn import NormalParamExtractor, TensorDictModule as Mod, TensorDictSequential as Seq, ProbabilisticTensorDictModule as Prob, CompositeDistribution
from tensordict import TensorDict
from torch import nn
num_cells = 32
device = "cpu"
num_joints = 4
num_wheels = 3

actor_net_base = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
)
actor_net_joints = nn.Sequential(
    nn.LazyLinear(num_cells, 2*num_joints, device=device),
    NormalParamExtractor(),
)
actor_net_wheels = nn.Sequential(
    nn.LazyLinear(num_cells, 2*num_wheels, device=device),
    NormalParamExtractor(),
)
actor_module = Seq(
    Mod(actor_net_base, in_keys=["obs"], out_keys=["embed"]),
    Mod(actor_net_joints, in_keys=["embed"], out_keys=[("params", "joints", "loc"), ("params", "joints", "scale")]),
    Mod(actor_net_joints, in_keys=["embed"], out_keys=[("params", "wheels", "loc"), ("params", "wheels", "scale")]),
)

# Define the policy module

actor = Seq(
    actor_module,
    Prob(
    in_keys=["params"],
    out_keys=["joints", "wheels"],
    distribution_class=CompositeDistribution,
    distribution_kwargs={
        "distribution_map": {"joints": d.Normal, "wheels": d.Normal},
    },
)
)
# With a tensordict as input
print(actor(TensorDict(obs=torch.randn(1, 3))))
# With a tensor as input
print(actor(obs=torch.randn(1, 3)))
</code></pre>
<p>You need to put the params in a container (e.g., “params” in our case) because otherwise you’ll override the <code>"joints"</code> and <code>"wheels"</code> entries during sampling.</p>

Mox_Antia: <p>Thank you for this solution and explanation!<br>
It is bit more clear now.<br>
So I have implemented the ProbabilisticActor as follow:</p>
<pre><code class="lang-auto">actor = ProbabilisticActor(
    module=actor_module,
    spec=env.action_spec,
    in_keys=["params"],
    out_keys=["joints", "wheels"],
    distribution_class=CompositeDistribution,
    distribution_kwargs={
        "distribution_map": {"joints": d.Normal, "wheels": d.Normal},
    },
    return_log_prob=True,
)
</code></pre>
<p>It works fine but the key “action” is missing which I believe is needed for the PPO training. I tried to add this line <code>"name_map": {"joints": "action", "wheels": "action"},</code>. But I get an unexpected keyword argument error on the CompositeDistribution.</p>
<p>Do you have a suggestion to solve this issue?</p>
<p>Thank you for all the help!</p>

vmoens: <p>try this</p>
<pre data-code-wrap="python"><code class="lang-python">policy = TensorDictSequential(
   policy,
   TensorDictModule(lambda x, y: torch.cat([x, y], -1), in_keys=["joints", "wheels"], out_keys=["action"])
)
</code></pre>

Mox_Antia: <aside class="quote no-group" data-username="vmoens" data-post="6" data-topic="210114">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vmoens/48/26121_2.png" class="avatar"> vmoens:</div>
<blockquote>
<pre><code class="lang-auto">policy = TensorDictSequential(
   policy,
   TensorDictModule(lambda x, y: torch.cat([x, y], -1), in_keys=["joints", "wheels"], out_keys=["action"])
)
</code></pre>
</blockquote>
</aside>
<p>It does not work because the distribution is then not accessible anymore.<br>
I get this error:</p>
<pre><code class="lang-auto">AttributeError: 'TensorDictSequential' object has no attribute 'get_dist
</code></pre>

vmoens: <p>Where are you getting that error?</p>

Mox_Antia: <p>I get this error at this line : <code>loss_vals = loss_module(subdata.to(device))</code></p>
<p>My code is the same as the PPO tutorial of TorchRL.</p>

vmoens: <p>Where is get_dist called?<br>
Can you share more info, like the full error stack and the lines of code causing it?</p>

Mox_Antia: <p>Yes sure,<br>
Here is the full error message:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1242, in __getattr__
    return super().__getattr__(name)
  File "/home/leni/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1688, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'TensorDictSequential' object has no attribute 'get_dist'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/leni/git/pyARE/ppo.py", line 252, in &lt;module&gt;
    loss_vals = loss_module(subdata.to(device))
  File "/home/leni/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1561, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/_contextlib.py", line 126, in decorate_context
    return func(*args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/nn/common.py", line 291, in wrapper
    return func(_self, tensordict, *args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/objectives/ppo.py", line 762, in forward
    log_weight, dist = self._log_weight(tensordict)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/objectives/ppo.py", line 474, in _log_weight
    dist = self.actor_network.get_dist(tensordict)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1249, in __getattr__
    raise err2 from err1
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/nn/common.py", line 1247, in __getattr__
    return getattr(super().__getattr__("module"), name)
  File "/home/leni/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1688, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'ModuleList' object has no attribute 'get_dist'
</code></pre>

vmoens: <p>Let’s backtrack for a second.<br>
If you need a single action vector, why building two separate dists? Why not just concatenating the two loc and scale vectors?<br>
We usually use two separate distributions if the actor expects two different components in the action but here my understanding is that you only need one, am I right?</p>

Mox_Antia: <p>Originally, I need two action vectors because I have two types of actuators: joints and wheels.<br>
They don’t have the same bounds. Also, wheels are controlled on velocity while the joints are controlled on position.</p>
<p>Though, I could use one action vector with normalised values.</p>

vmoens: <p>In TorchRL, the TanhNormal distribution accepts bounds that are not homogeneous. You can pass a vector <code>distribution_kwargs={"low": torch.tensor([-2.5, -2.5, -1.0, -1.0]), "high": torch.tensor([2.5, 2.5, 1.0, 1.0])}</code> and it will work just fine!</p>

Mox_Antia: <p>Ok, so, I modified the code to concatenate my joints and wheels in one loc and scale:</p>
<pre><code class="lang-auto">actor_module = TensorDictSequential(
    TensorDictModule(actor_net_base,in_keys=["observation"],out_keys=["embed"]),
    TensorDictModule(actor_net_joints,in_keys=["embed"],out_keys=[("params", "joints", "loc"), ("params", "joints", "scale")]),
    TensorDictModule(actor_net_wheels,in_keys=["embed"],out_keys=[("params", "wheels", "loc"), ("params", "wheels", "scale")]),
    TensorDictModule(lambda x, y, z, w: (torch.cat([x, y]), torch.cat([z,w])), 
                     in_keys=[("params", "joints", "loc"), ("params", "wheels", "loc"),("params", "joints", "scale"), ("params", "wheels", "scale")], 
                     out_keys=["loc","scale"]),
)

actor = ProbabilisticActor(
    module = actor_module,
    spec=env.action_spec,
    in_keys=["loc","scale"],
    out_keys=["joints", "wheels"],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "min": torch.cat([env.action_spec["joints"].space.low, env.action_spec["wheels"].space.low]),
         "max": torch.cat([env.action_spec["joints"].space.high, env.action_spec["wheels"].space.high]),
    },
    return_log_prob=True,
)
</code></pre>
<p>I don’t know if this is the best way to concatenate tensors.</p>
<p>Then, when I call my collector I have the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):                         
  File "/home/leni/git/pyARE/ppo.py", line 245, in &lt;module&gt; 
    for i, tensordict_data in enumerate(collector):
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 952, in iterator
    tensordict_out = self.rollout()
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/_utils.py", line 469, in unpack_rref_and_invoke_functio
n                            
    return func(self, *args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/collectors/collectors.py", line 1069, in rollout
    env_output, env_next_output = self.env.step_and_maybe_reset(env_input)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/envs/common.py", line 2576, in step_and_maybe_reset
    tensordict = self.step(tensordict)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/envs/common.py", line 1409, in step
    next_tensordict = self._step(tensordict)
  File "/home/leni/.local/lib/python3.10/site-packages/torchrl/envs/transforms/transforms.py", line 738, in _step
    next_tensordict = self.base_env._step(tensordict_in)
  File "/home/leni/git/pyARE/rl_env.py", line 71, in _step
    wheel_actions = action["wheels"].tolist()
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/base.py", line 242, in __getitem__
    result = self._get_tuple(idx_unravel, NO_DEFAULT)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/_td.py", line 1629, in _get_tuple
    first = self._get_str(key[0], default)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/_td.py", line 1625, in _get_str
    return self._default_get(first_key, default)
  File "/home/leni/.local/lib/python3.10/site-packages/tensordict/base.py", line 2430, in _default_get
    raise KeyError(
KeyError: 'key "wheels" not found in TensorDict with keys [\'collector\', \'done\', \'embed\', \'joints\', \'loc\', \'
observation\', \'params\', \'sample_log_prob\', \'scale\', \'step_count\', \'terminated\']'
</code></pre>
<p>I don’t understand why I loose the “wheels” key.</p>
<p>The collector is defined as follow:</p>
<pre><code class="lang-auto">collector = SyncDataCollector(
    env,
    actor,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)
</code></pre>

Mox_Antia: <p>I have decided to drop trying to have a composite action spec. I switched to one action tensor for the actor. Then, I split the action tensor into two tensors before sending the command to the simulator.</p>
<p>It seems to work now.</p>
<p>Thank you so much for all your help! I understand better how torch works now.</p>

vmoens: <p>Yeah I think it’s the simplest option! Glad it works!</p>

Gogh: <pre><code class="lang-auto">policy_net = torch.nn.Sequential(
     MultiAgentMLP(
        n_agent_inputs=3,  # n_obs_per_agent
        n_agent_outputs=4,  # 2 * n_actions_per_agents
        n_agents=1,
        centralised=False,  # the policies are decentralised (ie each agent will act from its observation)
        share_params=share_parameters_policy,
        device=env.device,
        depth=2,
        num_cells=64,
        activation_class=torch.nn.Tanh,
    ),
    NormalParamExtractor()
    
)
cat_module = TensorDictModule(
        lambda x, f, z: torch.cat([x,f, z], dim=-1),
        in_keys=[("agents", "observation","x"), ("agents", "observation","f"),("agents", "observation","z")],
        out_keys=[("agents", "obs")],
    )
policy_module = TensorDictModule(
    policy_net,
    in_keys=[("agents", "obs")],
    # in_keys=(("agents","observation")),
    # out_keys=[("agents", "action")],
    out_keys=[
        ("agents","action", "loc"), ("agents","action", "scale"),
              ],
)
uncat_module = TensorDictModule(
        lambda x, z: (x[...,0],z[...,0],x[...,1],z[...,1]),
        in_keys=[("agents","action", "loc"), ("agents","action", "scale")],
        out_keys=[
            ("agents","action","params","u_xx", "loc"),
            ("agents","action","params","u_xx", "scale"),
            ("agents","action","params","u_zz", "loc"),
            ("agents","action","params","u_zz", "scale"),
                  ],
    )
policy = ProbabilisticActor(
    module=TensorDictSequential(cat_module, policy_module,uncat_module),
    # spec=env.action_spec,
    in_keys=[("agents","action","params")],
    out_keys=[('agents', 'action', 'u_xx'), ('agents', 'action', 'u_zz')],
    distribution_class=CompositeDistribution,
    distribution_kwargs={
        "distribution_map": {( 'u_xx'): distributions.Normal, ( 'u_zz'): distributions.Normal},
        'name_map': {'u_xx':('agents', 'action','u_x'),'u_zz':('agents', 'action','u_z')}
    },
    
    return_log_prob=True,
    log_prob_key=("agents", "sample_log_prob"),
)
</code></pre>
<p>I think this messy code might do sth similiar to what you wanted. Note that the last key of the in_keys for the ProbabilisticActor must be “params”.</p>