diff --git a/mo_gymnasium/envs/four_room/four_room.py b/mo_gymnasium/envs/four_room/four_room.py index 5aa5e233..0f5d41cc 100644 --- a/mo_gymnasium/envs/four_room/four_room.py +++ b/mo_gymnasium/envs/four_room/four_room.py @@ -179,6 +179,9 @@ def step(self, action): # can now move self.state = (s1, collected) + if self.render_mode == "human": + self.render() + # into a goal cell if s1 == self.goal: phi = np.ones(len(self.all_shapes), dtype=np.float32) @@ -206,9 +209,6 @@ def step(self, action): phi = self.features(old_state, action, self.state) return self.state_to_array(self.state), phi, terminated, False, {} - if self.render_mode == "human": - self.render() - # into an empty cell return ( self.state_to_array(self.state),