Method step

source/isaaclab/isaaclab/envs/utils/marl.py:98–128

(self, action: torch.Tensor)

            return obs, extras

        def step(self, action: torch.Tensor) -> VecEnvStepReturn:
            # split single-agent actions to build the multi-agent ones
            # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces
            index = 0
            _actions = {}
            for agent in self.env.possible_agents:
                delta = gym.spaces.flatdim(self.env.action_spaces[agent])
                _actions[agent] = action[:, index : index + delta]
                index += delta

            # step the environment
            obs, rewards, terminated, time_outs, extras = self.env.step(_actions)

            # use environment state as observation
            if self._state_as_observation:
                obs = {"policy": self.env.state()}
            # concatenate agents' observations
            # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces
            else:
                obs = {
                    "policy": torch.cat(
                        [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1
                    )
                }

            # process environment outputs to return single-agent data
            rewards = sum(rewards.values())
            terminated = math.prod(terminated.values()).to(dtype=torch.bool)
            time_outs = math.prod(time_outs.values()).to(dtype=torch.bool)

            return obs, rewards, terminated, time_outs, extras

        def render(self, recompute: bool = False) -> np.ndarray | None:
            return self.env.render(recompute)