diff --git a/compiler_opt/rl/env.py b/compiler_opt/rl/env.py
index 15246bf..1bd96af 100644
--- a/compiler_opt/rl/env.py
+++ b/compiler_opt/rl/env.py
@@ -15,6 +15,8 @@
 """Gymlike environment definition for MLGO."""
 
 from __future__ import annotations
+import dataclasses
+from enum import Enum
 
 import math
 import subprocess
@@ -23,38 +25,31 @@
 import io
 import os
 import tempfile
-from typing import Any, Generator, List, Optional, Tuple, Type
+from typing import Callable, Generator, List, Optional, Tuple, Type
 
 import numpy as np
 
 from compiler_opt.rl import corpus
 from compiler_opt.rl import log_reader
 
-OBS_T = Any
 
-OBS_KEY = 'obs'
-REWARD_KEY = 'reward'
-SCORE_POLICY_KEY = 'score_policy'
-SCORE_DEFAULT_KEY = 'score_default'
-CONTEXT_KEY = 'context'
-MODULE_NAME_KEY = 'module_name'
-OBS_ID_KEY = 'obs_id'
-STEP_TYPE_KEY = 'step_type'
+class StepType(Enum):
+  FIRST = 1
+  MID = 2
+  LAST = 3
 
-FIRST_STEP_STR = 'first'
-MID_STEP_STR = 'mid'
-LAST_STEP_STR = 'last'
 
-_TERMINAL_OBS = {
-    OBS_KEY: {},
-    REWARD_KEY: 0.0,
-    SCORE_POLICY_KEY: 0.0,
-    SCORE_DEFAULT_KEY: 0.0,
-    CONTEXT_KEY: '',
-    MODULE_NAME_KEY: '',
-    OBS_ID_KEY: -1,
-    STEP_TYPE_KEY: LAST_STEP_STR,
-}
+@dataclasses.dataclass
+class TimeStep:
+  obs: Optional[dict[str, np.NDArray]]
+  reward: Optional[dict[str, float]]
+  score_policy: Optional[dict[str, float]]
+  score_default: Optional[dict[str, float]]
+  context: Optional[str]
+  module_name: str
+  obs_id: Optional[int]
+  step_type: StepType
+
 
 _INTERACTIVE_PIPE_FILE_BASE = 'interactive-pipe-base'
 
@@ -119,7 +114,8 @@
   associated to the default-compiled binary.
   """
 
-  def __init__(self, proc, get_scores_fn, module_name):
+  def __init__(self, proc: subprocess.Popen,
+               get_scores_fn: Callable[[], dict[str, float]], module_name):
     self._proc = proc
     self._get_scores_fn = get_scores_fn
     self._module_name = module_name
@@ -134,8 +130,8 @@
 
   def __init__(
       self,
-      proc,
-      get_scores_fn,
+      proc: subprocess.Popen,
+      get_scores_fn: Callable[[], dict[str, float]],
       module_name: str,
       reader_pipe: io.BufferedReader,
       writer_pipe: io.BufferedWriter,
@@ -147,18 +143,26 @@
 
     self._is_first_obs = True
 
-    self._terminal_obs = _TERMINAL_OBS
-    self._terminal_obs[MODULE_NAME_KEY] = module_name
+    self._terminal_obs = TimeStep(
+        obs=None,
+        reward=None,
+        score_policy=None,
+        score_default=None,
+        context=None,
+        module_name=module_name,
+        obs_id=None,
+        step_type=StepType.LAST,
+    )
 
   def _running(self) -> bool:
     return self._proc.poll() is None
 
-  def get_observation(self) -> OBS_T:
+  def get_observation(self) -> TimeStep:
     if not self._running():
       return self._terminal_obs
 
-    def _get_step_type():
-      step_type = FIRST_STEP_STR if self._is_first_obs else MID_STEP_STR
+    def _get_step_type() -> StepType:
+      step_type = StepType.FIRST if self._is_first_obs else StepType.MID
       self._is_first_obs = False
       return step_type
 
@@ -169,16 +173,16 @@
       for fv in obs.feature_values:
         array = fv.to_numpy()
         tv_dict[fv.spec.name] = np.reshape(array, newshape=fv.spec.shape)
-      return {
-          OBS_KEY: tv_dict,
-          REWARD_KEY: obs.score if obs.score else 0.0,
-          SCORE_POLICY_KEY: 0.0,
-          SCORE_DEFAULT_KEY: 0.0,
-          CONTEXT_KEY: obs.context,
-          MODULE_NAME_KEY: self._module_name,
-          OBS_ID_KEY: obs.observation_id,
-          STEP_TYPE_KEY: _get_step_type(),
-      }
+      return TimeStep(
+          obs=tv_dict,
+          reward={obs.context: obs.score} if obs.score else None,
+          score_policy=None,
+          score_default=None,
+          context=obs.context,
+          module_name=self._module_name,
+          obs_id=obs.observation_id,
+          step_type=_get_step_type(),
+      )
     except StopIteration:
       return self._terminal_obs
 
@@ -338,14 +342,14 @@
   def observation(self):
     return self._last_obs
 
-  def _get_observation(self) -> OBS_T:
+  def _get_observation(self) -> TimeStep:
     self._last_obs = self._iclang.get_observation()
-    if self._last_obs[STEP_TYPE_KEY] == 'last':
-      self._last_obs[SCORE_POLICY_KEY] = self._iclang.get_scores()
-      self._last_obs[SCORE_DEFAULT_KEY] = self._clang.get_scores()
-      self._last_obs[REWARD_KEY] = compute_relative_rewards(
-          self._last_obs[SCORE_POLICY_KEY], self._last_obs[SCORE_DEFAULT_KEY])
-    return self.observation()
+    if self._last_obs.step_type == StepType.LAST:
+      self._last_obs.score_policy = self._iclang.get_scores()
+      self._last_obs.score_default = self._clang.get_scores()
+      self._last_obs.reward = compute_relative_rewards(
+          self._last_obs.score_policy, self._last_obs.score_default)
+    return self._last_obs
 
   def reset(self, module: corpus.LoadedModuleSpec):
     # On the first call to reset(...), sending None starts the coroutine.
diff --git a/compiler_opt/rl/env_test.py b/compiler_opt/rl/env_test.py
index 1af508c..87577b3 100644
--- a/compiler_opt/rl/env_test.py
+++ b/compiler_opt/rl/env_test.py
@@ -155,10 +155,10 @@
       for idx in range(_NUM_STEPS):
         obs = clang_session.get_observation()
         self.assertEqual(
-            obs[env.OBS_KEY]['times_called'],
+            obs.obs['times_called'],
             np.array([idx], dtype=np.int64),
         )
-        self.assertEqual(obs[env.CONTEXT_KEY], f'context_{idx}')
+        self.assertEqual(obs.context, f'context_{idx}')
       mock_popen.assert_called_once()
 
 
@@ -178,15 +178,15 @@
     for env_itr in range(3):
       del env_itr
       step = test_env.reset(_MOCK_MODULE)
-      self.assertEqual(step[env.STEP_TYPE_KEY], env.FIRST_STEP_STR)
+      self.assertEqual(step.step_type, env.StepType.FIRST)
 
       for step_itr in range(_NUM_STEPS - 1):
         del step_itr
         step = test_env.step(np.array([1], dtype=np.int64))
-        self.assertEqual(step[env.STEP_TYPE_KEY], env.MID_STEP_STR)
+        self.assertEqual(step.step_type, env.StepType.MID)
 
       step = test_env.step(np.array([1], dtype=np.int64))
-      self.assertEqual(step[env.STEP_TYPE_KEY], env.LAST_STEP_STR)
+      self.assertEqual(step.step_type, env.StepType.LAST)
 
 
 if __name__ == '__main__':