[Doc] More depth in VMAS docs

torchrl/envs/libs/vmas.py

@@ -115,13 +115,21 @@ class VmasWrapper(_EnvWrapper):
         env (``vmas.simulator.environment.environment.Environment``): the vmas environment to wrap.
 
     Keyword Args:
-        num_envs (int): Number of vectorized simulation environments.
-        device (torch.device, optional): Device for simulation.
-        continuous_actions (bool, optional): Weather to use continuous actions. Defaults to ``True``.
-        max_steps (int, optional): Maximum number of steps in each vectorized environment after which done is returned.
-            Defaults to ``None`` (no truncation).
+        num_envs (int): Number of vectorized simulation environments. VMAS perfroms vectorized simulations using PyTorch.
+            This argument indicates the number of vectorized environments that should be simulated in a batch. It will also
+            determine the batch size of the environment.
+        device (torch.device, optional): Device for simulation. Defaults to the default device. All the tensors created by VMAS
+            will be placed on this device.
+        continuous_actions (bool, optional): Whether to use continuous actions. Defaults to ``True``. If ``False``, actions
+            will be discrete. The number of actions and their size will depend on the chosen scenario.
+            See the VMAS repository for more info.
+        max_steps (int, optional): Horizon of the task. Defaults to ``None`` (infinite horizon). Each VMAS scenario can
+            be terminating or not. If ``max_steps`` is specified,
+            the scenario is also terminated (and the ``"terminated"`` flag is set) whenever this horizon is reached.
+            Unlike gym's ``TimeLimit`` transform or torchrl's :class:`~torchrl.envs.transforms.StepCounter`,
+            this argument will not set the ``"truncated"`` entry in the tensordict.
         categorical_actions (bool, optional): if the environment actions are discrete, whether to transform
-            them to categorical or one-hot.
+            them to categorical or one-hot. Defaults to ``True``.
         group_map (MarlGroupMapType or Dict[str, List[str]], optional): how to group agents in tensordicts for
             input/output. By default, if the agent names follow the ``"<name>_<int>"``
             convention, they will be grouped by ``"<name>"``. If they do not follow this convention, they will be all put
@@ -627,22 +635,38 @@ class VmasEnv(VmasWrapper):
     Paper: https://arxiv.org/abs/2207.03530
 
     Args:
-        scenario (str or vmas.simulator.scenario.BaseScenario): the vmas environment to build.
+        scenario (str or vmas.simulator.scenario.BaseScenario): the vmas scenario to build.
+            Must be one of :attr:`~.available_envs`. For a description and rendering of available scenarios see
+            `the README <https://github.com/proroklab/VectorizedMultiAgentSimulator/tree/VMAS-1.3.3?tab=readme-ov-file#main-scenarios>`__.
+
 
     Keyword Args:
-        num_envs (int): Number of vectorized simulation environments.
-        device (torch.device, optional): Device for simulation.
-        continuous_actions (bool, optional): Weather to use continuous actions. Defaults to ``True``.
-        max_steps (int, optional): Maximum number of steps in each vectorized environment after which done is returned.
-            Defaults to ``None`` (no truncation).
+        num_envs (int): Number of vectorized simulation environments. VMAS perfroms vectorized simulations using PyTorch.
+            This argument indicates the number of vectorized environments that should be simulated in a batch. It will also
+            determine the batch size of the environment.
+        device (torch.device, optional): Device for simulation. Defaults to the defaultt device. All the tensors created by VMAS
+            will be placed on this device.
+        continuous_actions (bool, optional): Whether to use continuous actions. Defaults to ``True``. If ``False``, actions
+            will be discrete. The number of actions and their size will depend on the chosen scenario.
+            See the VMAS repositiory for more info.
+        max_steps (int, optional): Horizon of the task. Defaults to ``None`` (infinite horizon). Each VMAS scenario can
+            be terminating or not. If ``max_steps`` is specified,
+            the scenario is also terminated (and the ``"terminated"`` flag is set) whenever this horizon is reached.
+            Unlike gym's ``TimeLimit`` transform or torchrl's :class:`~torchrl.envs.transforms.StepCounter`,
+            this argument will not set the ``"truncated"`` entry in the tensordict.
         categorical_actions (bool, optional): if the environment actions are discrete, whether to transform
-            them to categorical or one-hot.
+            them to categorical or one-hot. Defaults to ``True``.
         group_map (MarlGroupMapType or Dict[str, List[str]], optional): how to group agents in tensordicts for
             input/output. By default, if the agent names follow the ``"<name>_<int>"``
             convention, they will be grouped by ``"<name>"``. If they do not follow this convention, they will be all put
             in one group named ``"agents"``.
             Otherwise, a group map can be specified or selected from some premade options.
             See :class:`~torchrl.envs.utils.MarlGroupMapType` for more info.
+        **kwargs (Dict, optional): These are additional arguments that can be passed to the VMAS scenario constructor.
+            (e.g., number of agents, reward sparsity). The available arguments will vary based on the chosen scenario.
+            To see the available arguments for a specific scenario, see the constructor in its file from
+            `the scenario folder <https://github.com/proroklab/VectorizedMultiAgentSimulator/tree/VMAS-1.3.3/vmas/scenarios>`__.
+
 
     Attributes:
         group_map (Dict[str, List[str]]): how to group agents in tensordicts for