add back achievement logging that was accidentally deleted

craftax/craftax/envs/common.py

@@ -0,0 +1,11 @@
+from craftax.craftax.craftax_state import EnvState
+from craftax.craftax.constants import *
+
+
+def log_achievements_to_info(state: EnvState, done: bool):
+    achievements = state.achievements * done * 100.0
+    info = {}
+    for achievement in Achievement:
+        name = f"Achievements/{achievement.name.lower()}"
+        info[name] = achievements[achievement.value]
+    return info

craftax/craftax/envs/craftax_pixels_env.py

@@ -4,6 +4,7 @@
 import chex
 
 from craftax.craftax.constants import *
+from craftax.craftax.envs.common import log_achievements_to_info
 from craftax.craftax.game_logic import craftax_step, is_game_over
 from craftax.craftax.craftax_state import EnvState, EnvParams, StaticEnvParams
 from craftax.craftax.renderer import render_craftax_pixels
@@ -34,7 +35,8 @@ def step_env(
         state, reward = craftax_step(key, state, action, params, self.static_env_params)
 
         done = self.is_terminal(state, params)
-        info = {"discount": self.discount(state, params)}
+        info = log_achievements_to_info(state, done)
+        info["discount"] = self.discount(state, params)
 
         return (
             lax.stop_gradient(self.get_obs(state)),
@@ -105,7 +107,8 @@ def step_env(
         state, reward = craftax_step(key, state, action, params, self.static_env_params)
 
         done = self.is_terminal(state, params)
-        info = {"discount": self.discount(state, params)}
+        info = log_achievements_to_info(state, done)
+        info["discount"] = self.discount(state, params)
 
         return (
             lax.stop_gradient(self.get_obs(state)),

craftax/craftax/envs/craftax_symbolic_env.py

@@ -4,6 +4,7 @@
 from typing import Tuple, Optional
 import chex
 
+from craftax.craftax.envs.common import log_achievements_to_info
 from craftax.environment_base.environment_bases import EnvironmentNoAutoReset
 from craftax.craftax.constants import *
 from craftax.craftax.game_logic import craftax_step, is_game_over
@@ -56,7 +57,8 @@ def step_env(
         state, reward = craftax_step(rng, state, action, params, self.static_env_params)
 
         done = self.is_terminal(state, params)
-        info = {"discount": self.discount(state, params)}
+        info = log_achievements_to_info(state, done)
+        info["discount"] = self.discount(state, params)
 
         return (
             lax.stop_gradient(self.get_obs(state)),
@@ -128,7 +130,8 @@ def step_env(
         state, reward = craftax_step(rng, state, action, params, self.static_env_params)
 
         done = self.is_terminal(state, params)
-        info = {"discount": self.discount(state, params)}
+        info = log_achievements_to_info(state, done)
+        info["discount"] = self.discount(state, params)
 
         return (
             lax.stop_gradient(self.get_obs(state)),