third_party/diffusion_policy/multirun_metrics.py

from typing import List, Optional
import pathlib
import pandas as pd
import numpy as np
import numba
import click
import time
import collections
import json
import wandb
import yaml
import numbers
import scipy.ndimage as sn
from diffusion_policy.common.json_logger import read_json_log, JsonLogger
import logging

@numba.jit(nopython=True)
def get_indexed_window_average(
        arr: np.ndarray, idxs: np.ndarray, window_size: int):
    result = np.zeros(idxs.shape, dtype=arr.dtype)
    length = arr.shape[0]
    for i in range(len(idxs)):
        idx = idxs[i]
        start = max(idx - window_size, 0)
        end = min(start + window_size, length)
        result[i] = np.mean(arr[start:end])
    return result


def compute_metrics(log_df: pd.DataFrame, key: str, 
        end_step: Optional[int]=None,
        k_min_loss: int=10,
        k_around_max: int=10,
        max_k_window: int=10,
        replace_slash: int=True,
        ):
    if key not in log_df:
        return dict()

    # prepare data
    if end_step is not None:
        log_df = log_df.iloc[:end_step]
    is_key = ~pd.isnull(log_df[key])
    is_key_idxs = is_key.index[is_key].to_numpy()
    if len(is_key_idxs) == 0:
        return dict()

    key_data = log_df[key][is_key].to_numpy()
    # after adding validation to workspace
    # rollout happens at the last step of each epoch
    # where the reported train_loss and val_loss
    # are already the average for that epoch
    train_loss = log_df['train_loss'][is_key].to_numpy()
    val_loss = log_df['val_loss'][is_key].to_numpy()

    result = dict()

    log_key = key
    if replace_slash:
        log_key = key.replace('/', '_')
    # max
    max_value = np.max(key_data)
    result['max/'+log_key] = max_value

    # k_around_max
    max_idx = np.argmax(key_data)
    end = min(max_idx + k_around_max // 2, len(key_data))
    start = max(end - k_around_max, 0)
    k_around_max_value = np.mean(key_data[start:end])
    result['k_around_max/'+log_key] = k_around_max_value

    # max_k_window
    k_window_value = sn.uniform_filter1d(key_data, size=max_k_window, axis=0, mode='nearest')
    max_k_window_value = np.max(k_window_value)
    result['max_k_window/'+log_key] = max_k_window_value

    # min_train_loss
    min_idx = np.argmin(train_loss)
    min_train_loss_value = key_data[min_idx]
    result['min_train_loss/'+log_key] = min_train_loss_value

    # min_val_loss
    min_idx = np.argmin(val_loss)
    min_val_loss_value = key_data[min_idx]
    result['min_val_loss/'+log_key] = min_val_loss_value

    # k_min_train_loss
    min_loss_idxs = np.argsort(train_loss)[:k_min_loss]
    k_min_train_loss_value = np.mean(key_data[min_loss_idxs])
    result['k_min_train_loss/'+log_key] = k_min_train_loss_value

    # k_min_val_loss
    min_loss_idxs = np.argsort(val_loss)[:k_min_loss]
    k_min_val_loss_value = np.mean(key_data[min_loss_idxs])
    result['k_min_val_loss/'+log_key] = k_min_val_loss_value

    # last
    result['last/'+log_key] = key_data[-1]

    # global step for visualization
    result['metric_global_step/'+log_key] = is_key_idxs[-1]
    return result


def compute_metrics_agg(
        log_dfs: List[pd.DataFrame], 
        key: str, end_step:int, 
        **kwargs):
    
    # compute metrics
    results = collections.defaultdict(list)
    for log_df in log_dfs:
        result = compute_metrics(log_df, key=key, end_step=end_step, **kwargs)
        for k, v in result.items():
            results[k].append(v)
    # agg
    agg_result = dict()
    for k, v in results.items():
        value = np.mean(v)
        if k.startswith('metric_global_step'):
            value = int(value)
        agg_result[k] = value
    return agg_result


@click.command()
@click.option('--input', '-i', required=True, help='Root logging dir, contains train_* dirs')
@click.option('--key', '-k', multiple=True, default=['test/mean_score'])
@click.option('--interval', default=10, type=float)
@click.option('--replace_slash', default=True, type=bool)
@click.option('--index_key', '-ik', default='global_step')
@click.option('--use_wandb', '-w', is_flag=True, default=False)
@click.option('--project', default=None)
@click.option('--name', default=None)
@click.option('--id', default=None)
@click.option('--group', default=None)
def main(
    input,
    key,
    interval,
    replace_slash,
    index_key,
    use_wandb,
    # wandb args
    project,
    name,
    id,
    group):
    root_dir = pathlib.Path(input)
    assert root_dir.is_dir()
    metrics_dir = root_dir.joinpath('metrics')
    metrics_dir.mkdir(exist_ok=True)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(str(metrics_dir.joinpath("metrics.log"))),
            logging.StreamHandler()
        ]
    )
    
    train_dirs = list(root_dir.glob('train_*'))
    log_files = [x.joinpath('logs.json.txt') for x in train_dirs]
    logging.info("Monitor waiting for log files!")
    while True:
        # wait for files to show up
        files_exist = True
        for log_file in log_files:
            if not log_file.is_file():
                files_exist = False
        if files_exist:
            break
        time.sleep(1.0)
    logging.info("All log files ready!")

    # init path
    metric_log_path = metrics_dir.joinpath('logs.json.txt')
    metric_path = metrics_dir.joinpath('metrics.json')
    config_path = root_dir.joinpath('config.yaml')

    # load config
    config = yaml.safe_load(config_path.open('r'))

    # init wandb
    wandb_run = None
    if use_wandb:
        wandb_kwargs = config['logging']
        if project is not None:
            wandb_kwargs['project'] = project
        if id is not None:
            wandb_kwargs['id'] = id
        if name is not None:
            wandb_kwargs['name'] = name
        if group is not None:
            wandb_kwargs['group'] = group
        wandb_kwargs['resume'] = True
        wandb_run = wandb.init(
            dir=str(metrics_dir),
            config=config,
            # auto-resume run, automatically load id
            # as long as using the same dir.
            # https://docs.wandb.ai/guides/track/advanced/resuming#resuming-guidance
            **wandb_kwargs
        )
        wandb.config.update(
            {
                "output_dir": str(root_dir),
            }
        )

    with JsonLogger(metric_log_path) as json_logger:
        last_log = json_logger.get_last_log()
        while True:
            # read json files
            log_dfs = [read_json_log(str(x), required_keys=key) for x in log_files]

            # previously logged data point
            last_log_idx = -1
            if last_log is not None:
                last_log_idx = log_dfs[0].index[log_dfs[0][index_key] <= last_log[index_key]][-1]
            
            start_idx = last_log_idx + 1
            # last idx where we have a data point from all logs
            end_idx = min(*[len(x) for x in log_dfs])

            # log every position
            for this_idx in range(start_idx, end_idx):
                # compute metrics
                all_metrics = dict()
                global_step = log_dfs[0]['global_step'][this_idx]
                epoch = log_dfs[0]['epoch'][this_idx]
                all_metrics['global_step'] = global_step
                all_metrics['epoch'] = epoch
                for k in key:
                    metrics = compute_metrics_agg(
                        log_dfs=log_dfs, key=k, end_step=this_idx+1,
                        replace_slash=replace_slash)
                    all_metrics.update(metrics)

                # sanitize metrics
                old_metrics = all_metrics
                all_metrics = dict()
                for k, v in old_metrics.items():
                    if isinstance(v, numbers.Integral):
                        all_metrics[k] = int(v)
                    elif isinstance(v, numbers.Number):
                        all_metrics[k] = float(v)
                
                has_update = all_metrics != last_log
                if has_update:
                    last_log = all_metrics
                    json_logger.log(all_metrics)

                    with metric_path.open('w') as f:
                        json.dump(all_metrics, f, sort_keys=True, indent=2)

                    if wandb_run is not None:
                        wandb_run.log(all_metrics, step=all_metrics[index_key])

                    logging.info(f"Metrics logged at step {all_metrics[index_key]}")
            
            time.sleep(interval)


if __name__ == "__main__":
    main()