bullet3/examples/pybullet/gym/pybullet_utils/logger.py

import pybullet_utils.mpi_util as MPIUtil

"""

Some simple logging functionality, inspired by rllab's logging.
Assumes that each diagnostic gets logged each iteration

Call logz.configure_output_file() to start logging to a 
tab-separated-values file (some_file_name.txt)

To load the learning curves, you can do, for example

A = np.genfromtxt('/tmp/expt_1468984536/log.txt',delimiter='\t',dtype=None, names=True)
A['EpRewMean']

"""

import os.path as osp, shutil, time, atexit, os, subprocess

class Logger:
    def print2(str):
        if (MPIUtil.is_root_proc()):
            print(str)
        return

    def __init__(self):
        self.output_file = None
        self.first_row = True
        self.log_headers = []
        self.log_current_row = {}
        self._dump_str_template = ""
        return

    def reset(self):
        self.first_row = True
        self.log_headers = []
        self.log_current_row = {}
        if self.output_file is not None:
            self.output_file = open(output_path, 'w')
        return

    def configure_output_file(self, filename=None):
        """
        Set output directory to d, or to /tmp/somerandomnumber if d is None
        """
        self.first_row = True
        self.log_headers = []
        self.log_current_row = {}

        output_path = filename or "output/log_%i.txt"%int(time.time())

        out_dir = os.path.dirname(output_path)
        if not os.path.exists(out_dir) and MPIUtil.is_root_proc():
            os.makedirs(out_dir)

        if (MPIUtil.is_root_proc()):
            self.output_file = open(output_path, 'w')
            assert osp.exists(output_path)
            atexit.register(self.output_file.close)

            Logger.print2("Logging data to " + self.output_file.name)
        return

    def log_tabular(self, key, val):
        """
        Log a value of some diagnostic
        Call this once for each diagnostic quantity, each iteration
        """
        if self.first_row and key not in self.log_headers:
            self.log_headers.append(key)
        else:
            assert key in self.log_headers, "Trying to introduce a new key %s that you didn't include in the first iteration"%key
        self.log_current_row[key] = val
        return

    def get_num_keys(self):
        return len(self.log_headers)

    def print_tabular(self):
        """
        Print all of the diagnostics from the current iteration
        """
        if (MPIUtil.is_root_proc()):
            vals = []
            Logger.print2("-"*37)
            for key in self.log_headers:
                val = self.log_current_row.get(key, "")
                if isinstance(val, float):
                    valstr = "%8.3g"%val
                elif isinstance(val, int):
                    valstr = str(val)
                else: 
                    valstr = val
                Logger.print2("| %15s | %15s |"%(key, valstr))
                vals.append(val)
            Logger.print2("-" * 37)
        return

    def dump_tabular(self):
        """
        Write all of the diagnostics from the current iteration
        """
        if (MPIUtil.is_root_proc()):
            if (self.first_row):
                self._dump_str_template = self._build_str_template()

            vals = []
            for key in self.log_headers:
                val = self.log_current_row.get(key, "")
                vals.append(val)
            
            if self.output_file is not None:
                if self.first_row:
                    header_str = self._dump_str_template.format(*self.log_headers)
                    self.output_file.write(header_str + "\n")

                val_str = self._dump_str_template.format(*map(str,vals))
                self.output_file.write(val_str + "\n")
                self.output_file.flush()

        self.log_current_row.clear()
        self.first_row=False
        return

    def _build_str_template(self):
        num_keys = self.get_num_keys()
        template = "{:<25}" * num_keys
        return template