def _step()

in robosumo/envs/sumo.py [0:0]

• 49 lines of code
• 18 McCabe index (conditional complexity)

    def _step(self, actions):
        if not self._mujoco_init:
            return self._get_obs(), 0, False, None

        dones = [False for _ in range(self._n_agents)]
        rewards = [0. for _ in range(self._n_agents)]
        infos = [{} for _ in range(self._n_agents)]

        # Call `before_step` on the agents
        for i in range(self._n_agents):
            self.agents[i].before_step()

        # Do simulation
        self.simulate(actions)

        # Call `after_step` on the agents
        for i in range(self._n_agents):
            infos[i]['ctrl_reward'] = self.agents[i].after_step(actions[i])

        # Get obs
        obs = self._get_obs()

        self._num_steps += 1

        # Compute rewards and dones
        for i, agent in enumerate(self.agents):
            self_xyz = agent.get_qpos()[:3]
            # Loose penalty
            infos[i]['lose_penalty'] = 0.
            if (self_xyz[2] < 0.29 or
                    np.max(np.abs(self_xyz[:2])) >= self._tatami_size):
                infos[i]['lose_penalty'] = - self.WIN_REWARD
                dones[i] = True
            # Win reward
            infos[i]['win_reward'] = 0.
            for opp in agent._opponents:
                opp_xyz = opp.get_qpos()[:3]
                if (opp_xyz[2] < 0.29 or
                        np.max(np.abs(opp_xyz[:2])) >= self._tatami_size):
                    infos[i]['win_reward'] += self.WIN_REWARD
                    infos[i]['winner'] = True
                    dones[i] = True
            infos[i]['main_reward'] = \
                infos[i]['win_reward'] + infos[i]['lose_penalty']
            # Draw penalty
            if self._num_steps > self._timestep_limit:
                infos[i]['main_reward'] += self.DRAW_PENALTY
                dones[i] = True
            # Move to opponent(s) and push them out of center
            infos[i]['move_to_opp_reward'] = 0.
            infos[i]['push_opp_reward'] = 0.
            for opp in agent._opponents:
                infos[i]['move_to_opp_reward'] += \
                    self._comp_move_reward(agent, opp.posafter)
                infos[i]['push_opp_reward'] += \
                    self._comp_push_reward(agent, opp.posafter)
            # Stay in center reward (unused)
            # infos[i]['stay_in_center'] = self._comp_stay_in_center_reward(agent)
            # Contact rewards and penalties (unused)
            # infos[i]['contact_reward'] = self._comp_contact_reward(agent)
            # Reward shaping
            infos[i]['shaping_reward'] = \
                infos[i]['ctrl_reward'] + \
                infos[i]['push_opp_reward'] + \
                infos[i]['move_to_opp_reward']
            # Add up rewards
            rewards[i] = infos[i]['main_reward'] + infos[i]['shaping_reward']

        rewards = tuple(rewards)
        dones = tuple(dones)
        infos = tuple(infos)

        return obs, rewards, dones, infos