Fixed problem that only affect in class visualizations

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