Batching a multicategorical spec

rsarpongstreetor: <p>Can I batch a multicategoral spec  as shown below. How do i it . I have tried inheriting the composite batch size but its not working. I need it to work in my  MARL environment .</p>
<p>if self.categorical_actions:<br>
self.action_spec = Composite(<br>
{<br>
(“agents”, “action”): MultiCategorical(<br>
nvec=torch.tensor([self.num_individual_actions] * self.obs_dim, device=self.device).repeat(self.num_agents, 1), # Repeat nvec for each agent<br>
shape=torch.Size([self.num_agents, self.obs_dim]), # Shape is (num_agents, obs_dim) per env<br>
dtype=torch.int64,<br>
device=self.device,<br>
)<br>
},</p>
<pre><code>         )
    else:
         self.action_spec = Composite(
             {
                 ("agents", "action"): UnboundedContinuousTensorSpec(
                     shape=torch.Size([self.num_agents, self.obs_dim]), # Shape is (num_agents, obs_dim) per env
                     dtype=torch.float32,
                     device=self.device,
                 )
             },

         )
</code></pre>

vmoens: <p>Thanks for the question</p>
<p>can you report:</p>
<ul>
<li>a complete example that I can reproduce easily</li>
<li>an error stack and an expected behaviour?</li>
</ul>

rsarpongstreetor: <p>I aimed to create and run a full BenchMARL experiment utilizing a custom VMAS scenario, a custom BenchMARL task, and custom policy and critic models. This involved several key stages: defining the custom scenario, defining the custom task, defining the custom models, configuring the experiment, and finally instantiating and running the experiment.</p>
<p>The process successfully defined the custom <code>MyScenario</code>, <code>Task</code>, <code>CustomPolicyModel</code>, and <code>CustomCriticModel</code> classes. However, significant challenges were encountered during the configuration and instantiation phases of the BenchMARL <code>Experiment</code>.</p>
<p>Multiple attempts were made to correctly structure the configuration dictionaries (<code>task_config</code>, <code>model_config</code>, <code>mappo_config_params</code>, <code>exp_config_for_experimentconfig</code>, <code>policy_model_config_structured</code>, <code>critic_model_config_structured</code>, <code>run_config</code>) and instantiate the necessary BenchMARL objects (<code>MappoConfig</code>, <code>ExperimentConfig</code>, <code>Experiment</code>). Debugging involved inspecting the signatures of BenchMARL classes (<code>MappoConfig.__init__</code>, <code>ExperimentConfig.__init__</code>, <code>Experiment.run</code>) to ensure parameters were correctly passed.</p>
<p><strong>The most critical and persistent issue arose when defining the action and value specifications using TorchRL’s <code>Composite</code> and <code>MultiCategorical</code> for the multi-agent discrete action space. Despite numerous attempts to adjust the shape and batch size arguments for these specs, a <code>ValueError</code> or <code>IndexError</code> related to shape mismatches consistently occurred during the instantiation of the <code>Experiment</code> object. This indicated a fundamental problem with how the scalar discrete action for each agent was being represented within the batched multi-agent specification structure expected by TorchRL/BenchMARL.</strong></p>
<h2><a name="p-473859-data-analysis-key-findings-1" class="anchor" href="#p-473859-data-analysis-key-findings-1"></a>Data Analysis Key Findings</h2>
<ul>
<li>
<p>The custom <code>MyScenario</code> class was successfully defined, including methods for data loading, environment setup, observation, reward, and termination based on a provided CSV data file.</p>
</li>
<li>
<p>The <code>Task</code> class inheriting from <code>BenchMARL.environments.common.TaskClass</code> was correctly implemented, providing necessary methods like <code>get_env_fun</code>, <code>get_specs</code>, <code>max_steps</code>, and <code>group_map</code>.</p>
</li>
<li>
<p>Custom policy (<code>CustomPolicyModel</code>) and critic (<code>CustomCriticModel</code>) network modules, including GNN layers and handling of batched multi-agent data, were successfully defined.</p>
</li>
<li>
<p><strong>The primary blocker was the inability to correctly define the <code>torchrl.data.Composite</code> specification for the multi-agent discrete action space using <code>torchrl.data.tensor_specs.MultiCategorical</code>. Attempts to define a single discrete action (3 options) per agent within the batched structure consistently resulted in <code>ValueError</code> or <code>IndexError</code> due to shape mismatches between the scalar action spec and the agent batch dimension.</strong></p>
</li>
<li>
<p>Similar issues were encountered when defining the value specification for the critic’s output using <code>Composite</code> and <code>Unbounded</code> specs.</p>
</li>
<li>
<p>Correctly partitioning configuration parameters between the <code>ExperimentConfig</code> object and the <code>Experiment.run</code> method required multiple iterations of inspecting class signatures and adjusting the configuration dictionaries.</p>
</li>
</ul>
<h2><a name="p-473859-insights-or-next-steps-2" class="anchor" href="#p-473859-insights-or-next-steps-2"></a>Insights or Next Steps</h2>
<ul>
<li>
<p><strong>The specific combination of TorchRL specs (<code>Composite</code> and <code>MultiCategorical</code>) for defining a multi-agent discrete action space with scalar actions per agent appears to have a complex or undocumented requirement regarding shape and batch size interaction that could not be resolved through standard interpretation of error messages. Further research or consultation of specific TorchRL/BenchMARL multi-agent discrete example implementations is needed.</strong></p>
</li>
<li>
<p>A simpler approach to defining the action space specification within the <code>Task.get_specs</code> method, potentially by relying more heavily on the environment’s own spec generation if possible, might circumvent the manual <code>Composite</code>/<code>MultiCategorical</code> construction issues.</p>
</li>
<li>
<p><strong>I have also used the pytorchrl difinitions for multi agents make_spec method using the combinations of specs and Composite and have faced similar issues above  with the ‘‘check env specs’’,</strong></p>
</li>
<li>
<p>Can you look into it the <code>Composite</code>/<code>MultiCategorical</code> and if possible the whole specs  construction issues.</p>
</li>
<li></li>
</ul>

vmoens: <p>Hello, as mentioned earlier an error stack or reprod example would be amazing!</p>

rsarpongstreetor: <p>I succeeded in implementing the MultiCategorical  by not batching the action spec as below, the batched  version was provided by  Envbase or Gemini</p>
<p>nvec=torch.full(( self.num_individual_actions_features,), self.num_individual_actions, dtype=torch.int64, device=self.device)</p>
<pre><code class="lang-auto">    agent_action_spec =MultiCategorical(nvec,
            shape=torch.Size([self.num_individual_actions_features,]),
            dtype=torch.int64,
            device=self.device
        )
    self.action_spec_unbatched = Composite(
          {("agents","action"): agent_action_spec}, batch_size=[self.num_agents,], device=self.device)

    self.action_spec = Composite(
         {("agents","action"): MultiCategorical(nvec, shape=torch.Size([*self.batch_size, self.num_individual_actions_features]), dtype=torch.int64, device=self.device)},
         batch_size=self.batch_size, device=self.device
    )
</code></pre>