<s> import json <EOL> import os <EOL> import time <EOL> from gym import error <EOL> class StatsRecorder ( object ) : <EOL> def __init__ ( self , directory , file_prefix ) : <EOL> self . initial_reset_timestamp = None <EOL> self . directory = directory <EOL> self . file_prefix = file_prefix <EOL> self . episode_lengths = [ ] <EOL> self . episode_rewards = [ ] <EOL> self . timestamps = [ ] <EOL> self . steps = None <EOL> self . rewards = None <EOL> self . done = None <EOL> def before_step ( self , action ) : <EOL> if self . done : <EOL> raise error . ResetNeeded ( "<STR_LIT>" ) <EOL> elif self . steps is None : <EOL> raise error . ResetNeeded ( "<STR_LIT>" ) <EOL> def after_step ( self , observation , reward , done , info ) : <EOL> self . steps += <NUM_LIT:1> <EOL> self . rewards += reward <EOL> if done : <EOL> self . done = True <EOL> def before_reset ( self ) : <EOL> self . done = False <EOL> if self . initial_reset_timestamp is None : <EOL> self . initial_reset_timestamp = time . time ( ) <EOL> def after_reset ( self , observation ) : <EOL> self . flush ( ) <EOL> def flush ( self ) : <EOL> if self . steps is not None : <EOL> self . episode_lengths . append ( self . steps ) <EOL> self . episode_rewards . append ( self . rewards ) <EOL> self . timestamps . append ( time . time ( ) ) <EOL> self . steps = <NUM_LIT:0> <EOL> self . rewards = <NUM_LIT:0> <EOL> def close ( self ) : <EOL> self . flush ( ) <EOL> filename = '<STR_LIT>' . format ( self . file_prefix , os . getpid ( ) ) <EOL> path = os . path . join ( self . directory , filename ) <EOL> with open ( path , '<STR_LIT:w>' ) as f : <EOL> json . dump ( { <EOL> '<STR_LIT>' : self . initial_reset_timestamp , <EOL> '<STR_LIT>' : self . timestamps , <EOL> '<STR_LIT>' : self . episode_lengths , <EOL> '<STR_LIT>' : self . episode_rewards , <EOL> } , f ) <EOL> return path </s>
