phantom/submodules/phantom-robomimic/robomimic/utils/hyperparam_utils.py

"""
A collection of utility functions and classes for generating config jsons for hyperparameter sweeps.
"""
import argparse
import os
import json
import re
import itertools

from collections import OrderedDict
from copy import deepcopy


class ConfigGenerator(object):
    """
    Useful class to keep track of hyperparameters to sweep, and to generate
    the json configs for each experiment run.
    """
    def __init__(self, base_config_file, wandb_proj_name="debug", script_file=None, generated_config_dir=None):
        """
        Args:
            base_config_file (str): path to a base json config to use as a starting point
                for the parameter sweep.

            script_file (str): script filename to write as output
        """
        assert isinstance(base_config_file, str)
        self.base_config_file = base_config_file
        assert generated_config_dir is None or isinstance(generated_config_dir, str)
        if generated_config_dir is not None:
            generated_config_dir = os.path.expanduser(generated_config_dir)
        self.generated_config_dir = generated_config_dir
        assert script_file is None or isinstance(script_file, str)
        if script_file is None:
            self.script_file = os.path.join('~', 'tmp/tmpp.sh')
        else:
            self.script_file = script_file
        self.script_file = os.path.expanduser(self.script_file)
        self.parameters = OrderedDict()

        assert isinstance(wandb_proj_name, str)
        self.wandb_proj_name = wandb_proj_name

    def add_param(self, key, name, group, values, value_names=None):
        """
        Add parameter to the hyperparameter sweep.

        Args:
            key (str): location of parameter in the config, using hierarchical key format
                (ex. train/data = config.train.data)

            name (str): name, as it will appear in the experiment name

            group (int): group id - parameters with the same ID have their values swept
                together

            values (list): list of values to sweep over for this parameter

            value_names ([str]): if provided, strings to use in experiment name for
                each value, instead of the parameter value. This is helpful for parameters
                that may have long or large values (for example, dataset path).
        """
        if value_names is not None:
            assert len(values) == len(value_names)
        self.parameters[key] = argparse.Namespace(
            key=key, 
            name=name, 
            group=group, 
            values=values, 
            value_names=value_names,
            hidename=hidename,
        )

    def generate(self):
        """
        Generates json configs for the hyperparameter sweep using attributes
        @self.parameters, @self.base_config_file, and @self.script_file,
        all of which should have first been set externally by calling
        @add_param, @set_base_config_file, and @set_script_file.
        """
        assert len(self.parameters) > 0, "must add parameters using add_param first!"
        generated_json_paths = self._generate_jsons()
        self._script_from_jsons(generated_json_paths)

    def _name_for_experiment(self, base_name, parameter_values, parameter_value_names):
        """
        This function generates the name for an experiment, given one specific
        parameter setting.

        Args:
            base_name (str): base experiment name
            parameter_values (OrderedDict): dictionary that maps parameter name to
                the parameter value for this experiment run
            parameter_value_names (dict): dictionary that maps parameter name to
                the name to use for its value in the experiment name

        Returns:
            name (str): generated experiment name
        """
        name = base_name
        for k in parameter_values:
            # append parameter name and value to end of base name
            if len(self.parameters[k].name) == 0 or self.parameters[k].hidename:
                # empty string indicates that naming should be skipped
                continue
            if len(self.parameters[k].name) == 0:
                # empty string indicates that naming should be skipped
                continue
            if parameter_value_names[k] is not None:
                # take name from passed dictionary
                val_str = parameter_value_names[k]
            else:
                val_str = parameter_values[k]
                if isinstance(parameter_values[k], list) or isinstance(parameter_values[k], tuple):
                    # convert list to string to avoid weird spaces and naming problems
                    val_str = "_".join([str(x) for x in parameter_values[k]])
            val_str = str(val_str)
            name += '_{}'.format(self.parameters[k].name)
            if len(val_str) > 0:
                name += '_{}'.format(val_str)
        return name

    def _get_parameter_ranges(self):
        """
        Extract parameter ranges from base json file. Also takes all possible
        combinations of the parameter ranges to generate an expanded set of values.

        Returns:
            parameter_ranges (dict): dictionary that maps the parameter to a list
                of all values it should take for each generated config. The length 
                of the list will be the total number of configs that will be
                generated from this scan.

            parameter_names (dict): dictionary that maps the parameter to a list
                of all name strings that should contribute to each invididual
                experiment's name. The length of the list will be the total 
                number of configs that will be generated from this scan.
        """

        # mapping from group id to list of indices to grab from each parameter's list 
        # of values in the parameter group
        parameter_group_indices = OrderedDict()
        for k in self.parameters:
            group_id = self.parameters[k].group
            assert isinstance(self.parameters[k].values, list)
            num_param_values = len(self.parameters[k].values)
            if group_id not in parameter_group_indices:
                parameter_group_indices[group_id] = list(range(num_param_values))
            else:
                assert len(parameter_group_indices[group_id]) == num_param_values, \
                    "error: inconsistent number of parameter values in group with id {}".format(group_id)

        keys = list(parameter_group_indices.keys())
        inds = list(parameter_group_indices.values())
        new_parameter_group_indices = OrderedDict(
            { k : [] for k in keys }
        )
        # get all combinations of the different parameter group indices
        # and then use these indices to determine the new parameter ranges
        # per member of each parameter group.
        #
        # e.g. with two parameter groups, one with two values, and another with three values
        # we have [0, 1] x [0, 1, 2] = [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]
        # so the corresponding parameter group indices are [0, 0, 0, 1, 1, 1] and 
        # [0, 1, 2, 0, 1, 2], and all parameters in each parameter group are indexed
        # together using these indices, to get each parameter range.
        for comb in itertools.product(*inds):
            for i in range(len(comb)):
                new_parameter_group_indices[keys[i]].append(comb[i])
        parameter_group_indices = new_parameter_group_indices

        # use the indices to gather the parameter values to sweep per parameter
        parameter_ranges = OrderedDict()
        parameter_names = OrderedDict()
        for k in self.parameters:
            parameter_values = self.parameters[k].values
            group_id = self.parameters[k].group
            inds = parameter_group_indices[group_id]
            parameter_ranges[k] = [parameter_values[ind] for ind in inds]

            # add in parameter names if supplied
            parameter_names[k] = None
            if self.parameters[k].value_names is not None:
                par_names = self.parameters[k].value_names
                assert isinstance(par_names, list)
                assert len(par_names) == len(parameter_values)
                parameter_names[k] = [par_names[ind] for ind in inds]

        # ensure that the number of parameter settings is the same per parameter
        first_key = list(parameter_ranges.keys())[0]
        num_settings = len(parameter_ranges[first_key])
        for k in parameter_ranges:
            assert len(parameter_ranges[k]) == num_settings, "inconsistent number of values"

        return parameter_ranges, parameter_names

    def _generate_jsons(self):
        """
        Generates json configs for the hyperparameter sweep, using @self.parameters and
        @self.base_config_file.

        Returns:
            json_paths (list): list of paths to created json files, one per experiment
        """

        # base directory for saving jsons
        if self.generated_config_dir:
            base_dir = self.generated_config_dir
            if not os.path.exists(base_dir):
                os.makedirs(base_dir)
        else:
            base_dir = os.path.abspath(os.path.dirname(self.base_config_file))

        # read base json
        base_config = load_json(self.base_config_file, verbose=False)

        # base exp name from this base config
        base_exp_name = base_config['experiment']['name']

        # use base json to determine the parameter ranges
        parameter_ranges, parameter_names = self._get_parameter_ranges()

        # iterate through each parameter setting to create each json
        first_key = list(parameter_ranges.keys())[0]
        num_settings = len(parameter_ranges[first_key])

        # keep track of path to generated jsons
        json_paths = []

        for i in range(num_settings):
            # the specific parameter setting for this experiment
            setting = { k : parameter_ranges[k][i] for k in parameter_ranges }
            maybe_parameter_names = OrderedDict()
            for k in parameter_names:
                maybe_parameter_names[k] = None
                if parameter_names[k] is not None:
                    maybe_parameter_names[k] = parameter_names[k][i]

            # experiment name from setting
            exp_name = self._name_for_experiment(
                base_name=base_exp_name, 
                parameter_values=setting, 
                parameter_value_names=maybe_parameter_names,
            )

            # copy old json, but override name, and parameter values
            json_dict = deepcopy(base_config)
            json_dict['experiment']['name'] = exp_name
            for k in parameter_ranges:
                set_value_for_key(json_dict, k, v=parameter_ranges[k][i])

            # populate list of identifying meta for logger;
            # see meta_config method in base_config.py for more info
            json_dict["experiment"]["logging"]["wandb_proj_name"] = self.wandb_proj_name
            if "meta" not in json_dict:
                json_dict["meta"] = dict()
            json_dict["meta"].update(
                hp_base_config_file=self.base_config_file,
                hp_keys=list(),
                hp_values=list(),
            )
            # logging: keep track of hyp param names and values as meta info
            for k in parameter_ranges.keys():
                key_name = self.parameters[k].name
                if key_name is not None and len(key_name) > 0:
                    if maybe_parameter_names[k] is not None:
                        value_name = maybe_parameter_names[k]
                    else:
                        value_name = setting[k]

                    json_dict["meta"]["hp_keys"].append(key_name)
                    json_dict["meta"]["hp_values"].append(value_name)

            # save file in same directory as old json
            json_path = os.path.join(base_dir, "{}.json".format(exp_name))
            save_json(json_dict, json_path)
            json_paths.append(json_path)

        print("Num exps:", len(json_paths))

        return json_paths

    def _script_from_jsons(self, json_paths):
        """
        Generates a bash script to run the experiments that correspond to
        the input jsons.
        """
        with open(self.script_file, 'w') as f:
            f.write("#!/bin/bash\n\n")
            for path in json_paths:
                # write python command to file
                cmd = "python train.py --config {}\n".format(path)

                print()
                print(cmd)
                f.write(cmd)


def load_json(json_file, verbose=True):
    """
    Simple utility function to load a json file as a dict.

    Args:
        json_file (str): path to json file to load
        verbose (bool): if True, pretty print the loaded json dictionary

    Returns:
        config (dict): json dictionary
    """
    with open(json_file, 'r') as f:
        config = json.load(f)
    if verbose:
        print('loading external config: =================')
        print(json.dumps(config, indent=4))
        print('==========================================')
    return config


def save_json(config, json_file):
    """
    Simple utility function to save a dictionary to a json file on disk.

    Args:
        config (dict): dictionary to save
        json_file (str): path to json file to write
    """
    with open(json_file, 'w') as f:
        # preserve original key ordering
        json.dump(config, f, sort_keys=False, indent=4)


def get_value_for_key(dic, k):
    """
    Get value for nested dictionary with levels denoted by "/" or ".".
    For example, if @k is "a/b", then this function returns
    @dic["a"]["b"].

    Args:
        dic (dict): a nested dictionary
        k (str): a single string meant to index several levels down into
            the nested dictionary, where levels can be denoted by "/" or
            by ".".
    Returns:
        val: the nested dictionary value for the provided key
    """
    val = dic
    subkeys = re.split('/|\.', k)
    for s in subkeys[:-1]:
        val = val[s]
    return val[subkeys[-1]]


def set_value_for_key(dic, k, v):
    """
    Set value for hierarchical dictionary with levels denoted by "/" or ".".

    Args:
        dic (dict): a nested dictionary
        k (str): a single string meant to index several levels down into
            the nested dictionary, where levels can be denoted by "/" or
            by ".".
        v: the value to set at the provided key
    """
    val = dic
    subkeys = re.split('/|\.', k) #k.split('/')
    for s in subkeys[:-1]:
        val = val[s]
    val[subkeys[-1]] = v