Make saving step metadata optional

This commit is contained in:
Marek Wolan
2023-12-04 10:42:20 +00:00
parent 8ea9db2d34
commit a5c4f7797d
5 changed files with 17 additions and 2 deletions

View File

@@ -40,7 +40,8 @@ class PrimaiteGymEnv(gymnasium.Env):
terminated = False
truncated = self.game.calculate_truncated()
info = {}
self._write_step_metadata_json(action, state, reward)
if self.game.save_step_metadata:
self._write_step_metadata_json(action, state, reward)
print(f"Episode: {self.game.episode_counter}, Step: {self.game.step_counter}, Reward: {reward}")
return next_obs, reward, terminated, truncated, info
@@ -183,7 +184,8 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
infos = {}
terminateds["__all__"] = len(self.terminateds) == len(self.agents)
truncateds["__all__"] = self.game.calculate_truncated()
self._write_step_metadata_json(actions, state, rewards)
if self.game.save_step_metadata:
self._write_step_metadata_json(actions, state, rewards)
return next_obs, rewards, terminateds, truncateds, infos
def _write_step_metadata_json(self, actions: Dict, state: Dict, rewards: Dict):