8c17debc9dc6552da8c317c0c852b55f03c80c61,slm_lab/env/__init__.py,Env,step,#Env#Any#,202

Before Change

        env_info_dict = self.u_env.step(action)
        reward = []
        state = []
        done = []
        for a, b in self.ab_proj:
            a_name = self.u_env.brain_names[a]
            a_env_info = env_info_dict[a_name]
            body_reward = a_env_info.rewards[b]
            reward.append(body_reward)
            body_state = a_env_info.states[b]
            state.append(body_state)
            body_done = a_env_info.local_done[b]
            done.append(body_done)
        return reward, state, done

    def close(self):
        self.u_env.close()

After Change

        action_e = util.flatten_nonan(action_e)
        env_info_dict = self.u_env.step(action_e)
        reward_e = self.data_spaces["reward"].init_data_s(e=self.e)
        state_e = self.data_spaces["state"].init_data_s(e=self.e)
        done_e = self.data_spaces["done"].init_data_s(e=self.e)
        for (a, b), body in util.ndenumerate_nonan(self.body_e):
            env_info_a = self.get_env_info(env_info_dict, a)
            reward_e[(a, b)] = env_info_a.rewards[b]
            state_e[(a, b)] = env_info_a.states[b]
            done_e[(a, b)] = env_info_a.local_done[b]
        return reward_e, state_e, done_e

Instances