commonml.helpers: Utility functions for array, environment, model operations

from .models import (

    clone

)

from .arrays import (

    homogenous_array,

    copy_tensor

)

from .env import (

    runner,

    rewards,

    get_from_env

)

 No newline at end of file

"""

Array or tensor operations operations.

"""

from typing import List, Iterable, Dict, Union

from collections import OrderedDict

import numpy as np

import torch

def homogenous_array(arrays: List[Iterable], start_align=True, fillvalue=np.nan) -> np.ndarray:

    """

    Convert a list of 1D arrays of multiple lengths into a 2D array padded with

    zeros.

    Parameters

    ----------

    arrays : List[Iterable]

        List of 1D iterables.

    start_align : bool, optional

        Whether to align all arrays' start positions, by default True

    fillvalue : float, int

        The value to put in empty parts of the 2D array, by default NaN

    Returns

    -------

    np.ndarray

        A 2D array of size len(arrays) x max array length

    """

    maxlen = max(map(len, arrays))

    res = np.zeros((len(arrays), maxlen)) + fillvalue

    for i, arr in enumerate(arrays):

        if start_align:

            res[i, :len(arr)] = arr

        else:

            res[i, -len(arr):] = arr

    return res

def copy_tensor(t: Union[torch.Tensor, Dict[str, torch.Tensor], List[torch.Tensor]]) \

    -> Union[torch.Tensor, Dict[str, torch.Tensor]]:

    """

    Make a copy of a tensor or a state_dict such that it is detached from the

    computation graph and does not share underlying data.

    Parameters

    ----------

    t : Union[torch.Tensor, Dict[str, torch.Tensor], List[torch.Tensor]]

        A tensor or a dictionary of [name, torch.Tensor]

    Returns

    -------

    Union[torch.Tensor, Dict[str, torch.Tensor]]

        Same object as t

    """

    if isinstance(t, OrderedDict):

        return OrderedDict([(k, v.clone().detach()) for k, v in t.items()])

    elif isinstance(t, torch.Tensor):

        return t.clone().detach()

    elif isinstance(t, (list, tuple)):

        return [t_.clone().detach() for t_ in t]

    else:

        raise TypeError('Only OrderedDict or Tensor supported')

"""

Environment utilities for gym classes.

"""

from typing import Callable, Dict, List, Tuple

import multiprocessing as mp

import numpy as np

import gym

def runner(env: gym.Env, callback: Callable, agent_fn: Callable=None,

           episodes: int=1, max_steps: int=1e3, processes: int=None):

    """

    Run environment and execute callback with local variables at every step.

    Parameters

    ----------

    env : gym.Env

        Environment to run.

    callback : Callable

        A function which accepts a dictionary of local variables. Called at each

        step of the environment.

    agent_fn : Callable, optional

        Function that accepts state and returns action to take for Env, by default random.

    episodes : int, optional

        Number of episodes to run for, by default 1

    max_steps : int, optional

        Maximum number of steps to take. Both `episodes` and `max_steps`

        will cause function to terminate if exceeded, by default 1e3

    processes : int, optional

        Not implemented. TODO.

    Returns

    -------

    List[np.ndarray]

        A list of arrays. Each array contains rewards for one episode.

    """

    ep, steps = 0, 0

    act = (lambda s: env.action_space.sample()) if agent_fn is None else agent_fn

    while (ep < episodes):

        done = False

        state = env.reset()

        while not done and (steps < max_steps):

            action = act(state)

            nstate, reward, done, _ = env.step(action)

            steps += 1

            callback(locals())

            state = nstate

        ep += 1

def rewards(env: gym.Env, agent_fn: Callable=None, episodes: int=1,

            max_steps: int=1e3) -> List[np.ndarray]:

    """

    Run environment and gather rewards.

    Parameters

    ----------

    env : gym.Env

        Environment to run.

    agent_fn : Callable, optional

        Function that accepts state and returns action to take for Env, by default random.

    episodes : int, optional

        Number of episodes to run for, by default 1

    max_steps : int, optional

        Maximum number of steps to take. Both `episodes` and `max_steps`

        will cause function to terminate if exceeded, by default 1e3

    Returns

    -------

    List[np.ndarray]

        A list of arrays. Each array contains rewards for one episode.

    """

    rewards = [[]]

    def rgetter(lcl: Dict):

        reward = lcl.get('reward')

        done = lcl.get('done')

        rewards[-1].append(reward)

        if done:

            rewards.append([])

    runner(env, rgetter, agent_fn, episodes=episodes, max_steps=max_steps)

    if len(rewards[-1]) == 0:

        rewards.pop()

    return [np.asarray(r) for r in rewards]

def get_from_env(variables: Tuple[str], env: gym.Env, agent_fn: Callable=None, episodes: int=1,

                 max_steps: int=1e3) -> Dict[str, List[np.ndarray]]:

    collection = {variable: [[]] for variable in variables}

    def getter(lcl: Dict):

        for variable in variables:

            value = lcl.get(variable)

            collection[variable][-1].append(value)

        done = lcl.get('done')

        if done:

            for variable in variables:

                collection[variable][-1] = np.asarray(collection[variable][-1])

                collection[variable].append([])

    runner(env, getter, agent_fn, episodes=episodes, max_steps=max_steps)

    for _, list_of_arrs in collection.items():

        if len(list_of_arrs[-1]) == 0:

            list_of_arrs.pop()

    return collection

 No newline at end of file

from typing import Any, Union, Dict

from copy import deepcopy

from torch.nn import Module

from sklearn.base import BaseEstimator

def clone(model: Union[BaseEstimator, Module], attrs: Dict[str, Any]=None) -> Union[BaseEstimator, Module]:

    """

    Copy a scikit-learn or pytorch model.

    Parameters

    ----------

    model : Union[BaseEstimator, Module]

        The model instance.

    attrs : Dict[str, Any], optional

        Any attributes to set in the copied model, by default None

    Returns

    -------

    Union[BaseEstimator, Module]

        The copied model with attributes set.

    """

    mcopy = deepcopy(model)

    if attrs is not None:

        for attr, val in attrs.items():

            setattr(mcopy, attr, val)

    return mcopy