in gym_recording/recording.py [0:0]
def save_complete(self):
"""
Save the latest batch and write a manifest listing all the batches.
We save the arrays as raw binary, in a format compatible with np.load.
We could possibly use numpy's compressed format, but the large observations we care about (VNC screens)
don't compress much, only by 30%, and it's a goal to be able to read the files from C++ or a browser someday.
"""
batch_fn = '{}.ep{:09}.json'.format(self.file_prefix, self.episodes_first)
bin_fn = '{}.ep{:09}.bin'.format(self.file_prefix, self.episodes_first)
with atomic_write.atomic_write(os.path.join(self.directory, batch_fn), False) as batch_f:
with atomic_write.atomic_write(os.path.join(self.directory, bin_fn), True) as bin_f:
def json_encode(obj):
if isinstance(obj, np.ndarray):
offset = bin_f.tell()
while offset%8 != 0:
bin_f.write(b'\x00')
offset += 1
obj.tofile(bin_f)
size = bin_f.tell() - offset
return {'__type': 'ndarray', 'shape': obj.shape, 'order': 'C', 'dtype': str(obj.dtype), 'npyfile': bin_fn, 'npyoff': offset, 'size': size}
return obj
json.dump({'episodes': self.episodes}, batch_f, default=json_encode)
bytes_per_step = float(bin_f.tell() + batch_f.tell()) / float(self.buffered_step_count)
self.batches.append({
'first': self.episodes_first,
'len': len(self.episodes),
'fn': batch_fn})
manifest = {'batches': self.batches}
manifest_fn = os.path.join(self.directory, '{}.manifest.json'.format(self.file_prefix))
with atomic_write.atomic_write(os.path.join(self.directory, manifest_fn), False) as f:
json.dump(manifest, f)
# Adjust batch size, aiming for 5 MB per file.
# This seems like a reasonable tradeoff between:
# writing speed (not too much overhead creating small files)
# local memory usage (buffering an entire batch before writing)
# random read access (loading the whole file isn't too much work when just grabbing one episode)
self.buffer_batch_size = max(1, min(50000, int(5000000 / bytes_per_step + 1)))
self.episodes = []
self.episodes_first = None
self.buffered_step_count = 0