diff --git a/irlc/ex01/agent.py b/irlc/ex01/agent.py
index 6cca61aa985b22c27d8f17b8210c3e8056e04e1a..1b277481104a8f421b3f821ae9403a3dedf77a26 100644
--- a/irlc/ex01/agent.py
+++ b/irlc/ex01/agent.py
@@ -306,7 +306,7 @@ def train(env,
                 if done or steps >= max_steps:
                     trajectory.state.append(sp)
                     trajectory.env_info.append(info_sp)
-                    trajectory.time.append(np.asarray(info_sp['time_seconds'] if 'time_seconds' in info_s else steps))
+                    trajectory.time.append(np.asarray(info_sp.get('time_seconds', None) if 'time_seconds' in info_s else steps))
                     break_outer = steps >= max_steps
                     break
                 s = sp