def save_act(self, path=None):
        """Save model to a pickle located at `path`"""
        if path is None:
            path = os.path.join(logger.get_dir(), "model.pkl")

        with tempfile.TemporaryDirectory() as td:
            save_variables(os.path.join(td, "model"))
            arc_name = os.path.join(td, "packed.zip")
            with zipfile.ZipFile(arc_name, 'w') as zipf:
                for root, dirs, files in os.walk(td):
                    for fname in files:
                        file_path = os.path.join(root, fname)
                        if file_path != arc_name:
                            zipf.write(file_path, os.path.relpath(file_path, td))
            with open(arc_name, "rb") as f:
                model_data = f.read()
        with open(path, "wb") as f:
            cloudpickle.dump((model_data, self._act_params), f)

def save_act ( self , path = None ) : if path is None : path = os . path . join ( logger . get_dir ( ) , "model.pkl" ) with tempfile . TemporaryDirectory ( ) as td : save_variables ( os . path . join ( td , "model" ) ) arc_name = os . path . join ( td , "packed.zip" ) with zipfile . ZipFile ( arc_name , 'w' ) as zipf : for root , dirs , files in os . walk ( td ) : for fname in files : file_path = os . path . join ( root , fname ) if file_path != arc_name : zipf . write ( file_path , os . path . relpath ( file_path , td ) ) with open ( arc_name , "rb" ) as f : model_data = f . read ( ) with open ( path , "wb" ) as f : cloudpickle . dump ( ( model_data , self . _act_params ) , f )


def nature_cnn(unscaled_images, **conv_kwargs):
    """
    CNN from Nature paper.
    """
    scaled_images = tf.cast(unscaled_images, tf.float32) / 255.
    activ = tf.nn.relu
    h = activ(conv(scaled_images, 'c1', nf=32, rf=8, stride=4, init_scale=np.sqrt(2),
                   **conv_kwargs))
    h2 = activ(conv(h, 'c2', nf=64, rf=4, stride=2, init_scale=np.sqrt(2), **conv_kwargs))
    h3 = activ(conv(h2, 'c3', nf=64, rf=3, stride=1, init_scale=np.sqrt(2), **conv_kwargs))
    h3 = conv_to_fc(h3)
    return activ(fc(h3, 'fc1', nh=512, init_scale=np.sqrt(2)))

def nature_cnn ( unscaled_images , ** conv_kwargs ) : scaled_images = tf . cast ( unscaled_images , tf . float32 ) / 255. activ = tf . nn . relu h = activ ( conv ( scaled_images , 'c1' , nf = 32 , rf = 8 , stride = 4 , init_scale = np . sqrt ( 2 ) , ** conv_kwargs ) ) h2 = activ ( conv ( h , 'c2' , nf = 64 , rf = 4 , stride = 2 , init_scale = np . sqrt ( 2 ) , ** conv_kwargs ) ) h3 = activ ( conv ( h2 , 'c3' , nf = 64 , rf = 3 , stride = 1 , init_scale = np . sqrt ( 2 ) , ** conv_kwargs ) ) h3 = conv_to_fc ( h3 ) return activ ( fc ( h3 , 'fc1' , nh = 512 , init_scale = np . sqrt ( 2 ) ) )


def conv_only(convs=[(32, 8, 4), (64, 4, 2), (64, 3, 1)], **conv_kwargs):
    '''
    convolutions-only net

    Parameters:
    ----------

    conv:       list of triples (filter_number, filter_size, stride) specifying parameters for each layer.

    Returns:

    function that takes tensorflow tensor as input and returns the output of the last convolutional layer

    '''

    def network_fn(X):
        out = tf.cast(X, tf.float32) / 255.
        with tf.variable_scope("convnet"):
            for num_outputs, kernel_size, stride in convs:
                out = layers.convolution2d(out,
                                           num_outputs=num_outputs,
                                           kernel_size=kernel_size,
                                           stride=stride,
                                           activation_fn=tf.nn.relu,
                                           **conv_kwargs)

        return out
    return network_fn

def conv_only ( convs = [ ( 32 , 8 , 4 ) , ( 64 , 4 , 2 ) , ( 64 , 3 , 1 ) ] , ** conv_kwargs ) : def network_fn ( X ) : out = tf . cast ( X , tf . float32 ) / 255. with tf . variable_scope ( "convnet" ) : for num_outputs , kernel_size , stride in convs : out = layers . convolution2d ( out , num_outputs = num_outputs , kernel_size = kernel_size , stride = stride , activation_fn = tf . nn . relu , ** conv_kwargs ) return out return network_fn


def make_vec_env(env_id, env_type, num_env, seed,
                 wrapper_kwargs=None,
                 start_index=0,
                 reward_scale=1.0,
                 flatten_dict_observations=True,
                 gamestate=None):
    """
    Create a wrapped, monitored SubprocVecEnv for Atari and MuJoCo.
    """
    wrapper_kwargs = wrapper_kwargs or {}
    mpi_rank = MPI.COMM_WORLD.Get_rank() if MPI else 0
    seed = seed + 10000 * mpi_rank if seed is not None else None
    logger_dir = logger.get_dir()
    def make_thunk(rank):
        return lambda: make_env(
            env_id=env_id,
            env_type=env_type,
            mpi_rank=mpi_rank,
            subrank=rank,
            seed=seed,
            reward_scale=reward_scale,
            gamestate=gamestate,
            flatten_dict_observations=flatten_dict_observations,
            wrapper_kwargs=wrapper_kwargs,
            logger_dir=logger_dir
        )

    set_global_seeds(seed)
    if num_env > 1:
        return SubprocVecEnv([make_thunk(i + start_index) for i in range(num_env)])
    else:
        return DummyVecEnv([make_thunk(start_index)])

def make_vec_env ( env_id , env_type , num_env , seed , wrapper_kwargs = None , start_index = 0 , reward_scale = 1.0 , flatten_dict_observations = True , gamestate = None ) : wrapper_kwargs = wrapper_kwargs or { } mpi_rank = MPI . COMM_WORLD . Get_rank ( ) if MPI else 0 seed = seed + 10000 * mpi_rank if seed is not None else None logger_dir = logger . get_dir ( ) def make_thunk ( rank ) : return lambda : make_env ( env_id = env_id , env_type = env_type , mpi_rank = mpi_rank , subrank = rank , seed = seed , reward_scale = reward_scale , gamestate = gamestate , flatten_dict_observations = flatten_dict_observations , wrapper_kwargs = wrapper_kwargs , logger_dir = logger_dir ) set_global_seeds ( seed ) if num_env > 1 : return SubprocVecEnv ( [ make_thunk ( i + start_index ) for i in range ( num_env ) ] ) else : return DummyVecEnv ( [ make_thunk ( start_index ) ] )

Create a wrapped monitored SubprocVecEnv for Atari and MuJoCo .


def parse_unknown_args(args):
    """
    Parse arguments not consumed by arg parser into a dicitonary
    """
    retval = {}
    preceded_by_key = False
    for arg in args:
        if arg.startswith('--'):
            if '=' in arg:
                key = arg.split('=')[0][2:]
                value = arg.split('=')[1]
                retval[key] = value
            else:
                key = arg[2:]
                preceded_by_key = True
        elif preceded_by_key:
            retval[key] = arg
            preceded_by_key = False

    return retval

def parse_unknown_args ( args ) : retval = { } preceded_by_key = False for arg in args : if arg . startswith ( '--' ) : if '=' in arg : key = arg . split ( '=' ) [ 0 ] [ 2 : ] value = arg . split ( '=' ) [ 1 ] retval [ key ] = value else : key = arg [ 2 : ] preceded_by_key = True elif preceded_by_key : retval [ key ] = arg preceded_by_key = False return retval


def clear_mpi_env_vars():
    """
    from mpi4py import MPI will call MPI_Init by default.  If the child process has MPI environment variables, MPI will think that the child process is an MPI process just like the parent and do bad things such as hang.
    This context manager is a hacky way to clear those environment variables temporarily such as when we are starting multiprocessing
    Processes.
    """
    removed_environment = {}
    for k, v in list(os.environ.items()):
        for prefix in ['OMPI_', 'PMI_']:
            if k.startswith(prefix):
                removed_environment[k] = v
                del os.environ[k]
    try:
        yield
    finally:
        os.environ.update(removed_environment)

def clear_mpi_env_vars ( ) : removed_environment = { } for k , v in list ( os . environ . items ( ) ) : for prefix in [ 'OMPI_' , 'PMI_' ] : if k . startswith ( prefix ) : removed_environment [ k ] = v del os . environ [ k ] try : yield finally : os . environ . update ( removed_environment )

from mpi4py import MPI will call MPI_Init by default . If the child process has MPI environment variables MPI will think that the child process is an MPI process just like the parent and do bad things such as hang . This context manager is a hacky way to clear those environment variables temporarily such as when we are starting multiprocessing Processes .


def cg(f_Ax, b, cg_iters=10, callback=None, verbose=False, residual_tol=1e-10):
    """
    Demmel p 312
    """
    p = b.copy()
    r = b.copy()
    x = np.zeros_like(b)
    rdotr = r.dot(r)

    fmtstr =  "%10i %10.3g %10.3g"
    titlestr =  "%10s %10s %10s"
    if verbose: print(titlestr % ("iter", "residual norm", "soln norm"))

    for i in range(cg_iters):
        if callback is not None:
            callback(x)
        if verbose: print(fmtstr % (i, rdotr, np.linalg.norm(x)))
        z = f_Ax(p)
        v = rdotr / p.dot(z)
        x += v*p
        r -= v*z
        newrdotr = r.dot(r)
        mu = newrdotr/rdotr
        p = r + mu*p

        rdotr = newrdotr
        if rdotr < residual_tol:
            break

    if callback is not None:
        callback(x)
    if verbose: print(fmtstr % (i+1, rdotr, np.linalg.norm(x)))  # pylint: disable=W0631
    return x

def
cg
(
f_Ax
,
b
,
cg_iters
=
10
,
callback
=
,
verbose
=
False
,
residual_tol
=
1e-10
)
:
p
=
b
.
copy
(
)
r
=
b
.
copy
(
)
x
=
np
.
zeros_like
(
b
)
rdotr
=
r
.
dot
(
r
)
fmtstr
=
"%10i %10.3g %10.3g"
titlestr
=
"%10s %10s %10s"
if
verbose
:
print
(
titlestr
%
(
"iter"
,
"residual norm"
,
"soln norm"
)
)
for
i
in
range
(
cg_iters
)
:
if
callback
is
not
:
callback
(
x
)
if
verbose
:
print
(
fmtstr
%
(
i
,
rdotr
,
np
.
linalg
.
norm
(
x
)
)
)
z
=
f_Ax
(
p
)
v
=
rdotr
/
p
.
dot
(
z
)
x
+=
v
*
p
r
-=
v
*
z
newrdotr
=
r
.
dot
(
r
)
mu
=
newrdotr
/
rdotr
p
=
r
+
mu
*
p
rdotr
=
newrdotr
if
rdotr
<
residual_tol
:
break
if
callback
is
not
:
callback
(
x
)
if
verbose
:
print
(
fmtstr
%
(
i
+
1
,
rdotr
,
np
.
linalg
.
norm
(
x
)
)
)
return
x


def observation_placeholder(ob_space, batch_size=None, name='Ob'):
    '''
    Create placeholder to feed observations into of the size appropriate to the observation space

    Parameters:
    ----------

    ob_space: gym.Space     observation space

    batch_size: int         size of the batch to be fed into input. Can be left None in most cases.

    name: str               name of the placeholder

    Returns:
    -------

    tensorflow placeholder tensor
    '''

    assert isinstance(ob_space, Discrete) or isinstance(ob_space, Box) or isinstance(ob_space, MultiDiscrete), \
        'Can only deal with Discrete and Box observation spaces for now'

    dtype = ob_space.dtype
    if dtype == np.int8:
        dtype = np.uint8

    return tf.placeholder(shape=(batch_size,) + ob_space.shape, dtype=dtype, name=name)


def observation_input(ob_space, batch_size=None, name='Ob'):
    '''
    Create placeholder to feed observations into of the size appropriate to the observation space, and add input
    encoder of the appropriate type.
    '''

    placeholder = observation_placeholder(ob_space, batch_size, name)
    return placeholder, encode_observation(ob_space, placeholder)

def observation_input ( ob_space , batch_size = None , name = 'Ob' ) : placeholder = observation_placeholder ( ob_space , batch_size , name ) return placeholder , encode_observation ( ob_space , placeholder )

Create placeholder to feed observations into of the size appropriate to the observation space and add input encoder of the appropriate type .


def encode_observation(ob_space, placeholder):
    '''
    Encode input in the way that is appropriate to the observation space

    Parameters:
    ----------

    ob_space: gym.Space             observation space

    placeholder: tf.placeholder     observation input placeholder
    '''
    if isinstance(ob_space, Discrete):
        return tf.to_float(tf.one_hot(placeholder, ob_space.n))
    elif isinstance(ob_space, Box):
        return tf.to_float(placeholder)
    elif isinstance(ob_space, MultiDiscrete):
        placeholder = tf.cast(placeholder, tf.int32)
        one_hots = [tf.to_float(tf.one_hot(placeholder[..., i], ob_space.nvec[i])) for i in range(placeholder.shape[-1])]
        return tf.concat(one_hots, axis=-1)
    else:
        raise NotImplementedError

def encode_observation ( ob_space , placeholder ) : if isinstance ( ob_space , Discrete ) : return tf . to_float ( tf . one_hot ( placeholder , ob_space . n ) ) elif isinstance ( ob_space , Box ) : return tf . to_float ( placeholder ) elif isinstance ( ob_space , MultiDiscrete ) : placeholder = tf . cast ( placeholder , tf . int32 ) one_hots = [ tf . to_float ( tf . one_hot ( placeholder [ ... , i ] , ob_space . nvec [ i ] ) ) for i in range ( placeholder . shape [ - 1 ] ) ] return tf . concat ( one_hots , axis = - 1 ) else : raise NotImplementedError


def save_policy(self, path):
        """Pickles the current policy for later inspection.
        """
        with open(path, 'wb') as f:
            pickle.dump(self.policy, f)

def save_policy ( self , path ) : with open ( path , 'wb' ) as f : pickle . dump ( self . policy , f )

Pickles the current policy for later inspection .


def logs(self, prefix='worker'):
        """Generates a dictionary that contains all collected statistics.
        """
        logs = []
        logs += [('success_rate', np.mean(self.success_history))]
        if self.compute_Q:
            logs += [('mean_Q', np.mean(self.Q_history))]
        logs += [('episode', self.n_episodes)]

        if prefix != '' and not prefix.endswith('/'):
            return [(prefix + '/' + key, val) for key, val in logs]
        else:
            return logs

def logs ( self , prefix = 'worker' ) : logs = [ ] logs += [ ( 'success_rate' , np . mean ( self . success_history ) ) ] if self . compute_Q : logs += [ ( 'mean_Q' , np . mean ( self . Q_history ) ) ] logs += [ ( 'episode' , self . n_episodes ) ] if prefix != '' and not prefix . endswith ( '/' ) : return [ ( prefix + '/' + key , val ) for key , val in logs ] else : return logs

Generates a dictionary that contains all collected statistics .


def smooth(y, radius, mode='two_sided', valid_only=False):
    '''
    Smooth signal y, where radius is determines the size of the window

    mode='twosided':
        average over the window [max(index - radius, 0), min(index + radius, len(y)-1)]
    mode='causal':
        average over the window [max(index - radius, 0), index]

    valid_only: put nan in entries where the full-sized window is not available

    '''
    assert mode in ('two_sided', 'causal')
    if len(y) < 2*radius+1:
        return np.ones_like(y) * y.mean()
    elif mode == 'two_sided':
        convkernel = np.ones(2 * radius+1)
        out = np.convolve(y, convkernel,mode='same') / np.convolve(np.ones_like(y), convkernel, mode='same')
        if valid_only:
            out[:radius] = out[-radius:] = np.nan
    elif mode == 'causal':
        convkernel = np.ones(radius)
        out = np.convolve(y, convkernel,mode='full') / np.convolve(np.ones_like(y), convkernel, mode='full')
        out = out[:-radius+1]
        if valid_only:
            out[:radius] = np.nan
    return out

def smooth ( y , radius , mode = 'two_sided' , valid_only = False ) : assert mode in ( 'two_sided' , 'causal' ) if len ( y ) < 2 * radius + 1 : return np . ones_like ( y ) * y . mean ( ) elif mode == 'two_sided' : convkernel = np . ones ( 2 * radius + 1 ) out = np . convolve ( y , convkernel , mode = 'same' ) / np . convolve ( np . ones_like ( y ) , convkernel , mode = 'same' ) if valid_only : out [ : radius ] = out [ - radius : ] = np . nan elif mode == 'causal' : convkernel = np . ones ( radius ) out = np . convolve ( y , convkernel , mode = 'full' ) / np . convolve ( np . ones_like ( y ) , convkernel , mode = 'full' ) out = out [ : - radius + 1 ] if valid_only : out [ : radius ] = np . nan return out

Smooth signal y where radius is determines the size of the window


def copy_obs_dict(obs):
    """
    Deep-copy an observation dict.
    """
    return {k: np.copy(v) for k, v in obs.items()}

def copy_obs_dict ( obs ) : return { k : np . copy ( v ) for k , v in obs . items ( ) }


def obs_space_info(obs_space):
    """
    Get dict-structured information about a gym.Space.

    Returns:
      A tuple (keys, shapes, dtypes):
        keys: a list of dict keys.
        shapes: a dict mapping keys to shapes.
        dtypes: a dict mapping keys to dtypes.
    """
    if isinstance(obs_space, gym.spaces.Dict):
        assert isinstance(obs_space.spaces, OrderedDict)
        subspaces = obs_space.spaces
    else:
        subspaces = {None: obs_space}
    keys = []
    shapes = {}
    dtypes = {}
    for key, box in subspaces.items():
        keys.append(key)
        shapes[key] = box.shape
        dtypes[key] = box.dtype
    return keys, shapes, dtypes

def obs_space_info ( obs_space ) : if isinstance ( obs_space , gym . spaces . Dict ) : assert isinstance ( obs_space . spaces , OrderedDict ) subspaces = obs_space . spaces else : subspaces = { None : obs_space } keys = [ ] shapes = { } dtypes = { } for key , box in subspaces . items ( ) : keys . append ( key ) shapes [ key ] = box . shape dtypes [ key ] = box . dtype return keys , shapes , dtypes

Get dict - structured information about a gym . Space .


def q_retrace(R, D, q_i, v, rho_i, nenvs, nsteps, gamma):
    """
    Calculates q_retrace targets

    :param R: Rewards
    :param D: Dones
    :param q_i: Q values for actions taken
    :param v: V values
    :param rho_i: Importance weight for each action
    :return: Q_retrace values
    """
    rho_bar = batch_to_seq(tf.minimum(1.0, rho_i), nenvs, nsteps, True)  # list of len steps, shape [nenvs]
    rs = batch_to_seq(R, nenvs, nsteps, True)  # list of len steps, shape [nenvs]
    ds = batch_to_seq(D, nenvs, nsteps, True)  # list of len steps, shape [nenvs]
    q_is = batch_to_seq(q_i, nenvs, nsteps, True)
    vs = batch_to_seq(v, nenvs, nsteps + 1, True)
    v_final = vs[-1]
    qret = v_final
    qrets = []
    for i in range(nsteps - 1, -1, -1):
        check_shape([qret, ds[i], rs[i], rho_bar[i], q_is[i], vs[i]], [[nenvs]] * 6)
        qret = rs[i] + gamma * qret * (1.0 - ds[i])
        qrets.append(qret)
        qret = (rho_bar[i] * (qret - q_is[i])) + vs[i]
    qrets = qrets[::-1]
    qret = seq_to_batch(qrets, flat=True)
    return qret

def q_retrace ( R , D , q_i , v , rho_i , nenvs , nsteps , gamma ) : rho_bar = batch_to_seq ( tf . minimum ( 1.0 , rho_i ) , nenvs , nsteps , True ) rs = batch_to_seq ( R , nenvs , nsteps , True ) ds = batch_to_seq ( D , nenvs , nsteps , True ) q_is = batch_to_seq ( q_i , nenvs , nsteps , True ) vs = batch_to_seq ( v , nenvs , nsteps + 1 , True ) v_final = vs [ - 1 ] qret = v_final qrets = [ ] for i in range ( nsteps - 1 , - 1 , - 1 ) : check_shape ( [ qret , ds [ i ] , rs [ i ] , rho_bar [ i ] , q_is [ i ] , vs [ i ] ] , [ [ nenvs ] ] * 6 ) qret = rs [ i ] + gamma * qret * ( 1.0 - ds [ i ] ) qrets . append ( qret ) qret = ( rho_bar [ i ] * ( qret - q_is [ i ] ) ) + vs [ i ] qrets = qrets [ : : - 1 ] qret = seq_to_batch ( qrets , flat = True ) return qret


def value(self, t):
        """See Schedule.value"""
        for (l_t, l), (r_t, r) in zip(self._endpoints[:-1], self._endpoints[1:]):
            if l_t <= t and t < r_t:
                alpha = float(t - l_t) / (r_t - l_t)
                return self._interpolation(l, r, alpha)

        # t does not belong to any of the pieces, so doom.
        assert self._outside_value is not None
        return self._outside_value

def value ( self , t ) : for ( l_t , l ) , ( r_t , r ) in zip ( self . _endpoints [ : - 1 ] , self . _endpoints [ 1 : ] ) : if l_t <= t and t < r_t : alpha = float ( t - l_t ) / ( r_t - l_t ) return self . _interpolation ( l , r , alpha ) assert self . _outside_value is not None return self . _outside_value


def _subproc_worker(pipe, parent_pipe, env_fn_wrapper, obs_bufs, obs_shapes, obs_dtypes, keys):
    """
    Control a single environment instance using IPC and
    shared memory.
    """
    def _write_obs(maybe_dict_obs):
        flatdict = obs_to_dict(maybe_dict_obs)
        for k in keys:
            dst = obs_bufs[k].get_obj()
            dst_np = np.frombuffer(dst, dtype=obs_dtypes[k]).reshape(obs_shapes[k])  # pylint: disable=W0212
            np.copyto(dst_np, flatdict[k])

    env = env_fn_wrapper.x()
    parent_pipe.close()
    try:
        while True:
            cmd, data = pipe.recv()
            if cmd == 'reset':
                pipe.send(_write_obs(env.reset()))
            elif cmd == 'step':
                obs, reward, done, info = env.step(data)
                if done:
                    obs = env.reset()
                pipe.send((_write_obs(obs), reward, done, info))
            elif cmd == 'render':
                pipe.send(env.render(mode='rgb_array'))
            elif cmd == 'close':
                pipe.send(None)
                break
            else:
                raise RuntimeError('Got unrecognized cmd %s' % cmd)
    except KeyboardInterrupt:
        print('ShmemVecEnv worker: got KeyboardInterrupt')
    finally:
        env.close()

Control a single environment instance using IPC and shared memory .


def learn(
    network,
    env,
    seed=None,
    nsteps=5,
    total_timesteps=int(80e6),
    vf_coef=0.5,
    ent_coef=0.01,
    max_grad_norm=0.5,
    lr=7e-4,
    lrschedule='linear',
    epsilon=1e-5,
    alpha=0.99,
    gamma=0.99,
    log_interval=100,
    load_path=None,
    **network_kwargs):

    '''
    Main entrypoint for A2C algorithm. Train a policy with given network architecture on a given environment using a2c algorithm.

    Parameters:
    -----------

    network:            policy network architecture. Either string (mlp, lstm, lnlstm, cnn_lstm, cnn, cnn_small, conv_only - see baselines.common/models.py for full list)
                        specifying the standard network architecture, or a function that takes tensorflow tensor as input and returns
                        tuple (output_tensor, extra_feed) where output tensor is the last network layer output, extra_feed is None for feed-forward
                        neural nets, and extra_feed is a dictionary describing how to feed state into the network for recurrent neural nets.
                        See baselines.common/policies.py/lstm for more details on using recurrent nets in policies

    env:                RL environment. Should implement interface similar to VecEnv (baselines.common/vec_env) or be wrapped with DummyVecEnv (baselines.common/vec_env/dummy_vec_env.py)

    seed:               seed to make random number sequence in the alorightm reproducible. By default is None which means seed from system noise generator (not reproducible)

    nsteps:             int, number of steps of the vectorized environment per update (i.e. batch size is nsteps * nenv where
                        nenv is number of environment copies simulated in parallel)

    total_timesteps:    int, total number of timesteps to train on (default: 80M)

    vf_coef:            float, coefficient in front of value function loss in the total loss function (default: 0.5)

    ent_coef:           float, coeffictiant in front of the policy entropy in the total loss function (default: 0.01)

    max_gradient_norm:  float, gradient is clipped to have global L2 norm no more than this value (default: 0.5)

    lr:                 float, learning rate for RMSProp (current implementation has RMSProp hardcoded in) (default: 7e-4)

    lrschedule:         schedule of learning rate. Can be 'linear', 'constant', or a function [0..1] -> [0..1] that takes fraction of the training progress as input and
                        returns fraction of the learning rate (specified as lr) as output

    epsilon:            float, RMSProp epsilon (stabilizes square root computation in denominator of RMSProp update) (default: 1e-5)

    alpha:              float, RMSProp decay parameter (default: 0.99)

    gamma:              float, reward discounting parameter (default: 0.99)

    log_interval:       int, specifies how frequently the logs are printed out (default: 100)

    **network_kwargs:   keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
                        For instance, 'mlp' network architecture has arguments num_hidden and num_layers.

    '''

    set_global_seeds(seed)

    # Get the nb of env
    nenvs = env.num_envs
    policy = build_policy(env, network, **network_kwargs)

    # Instantiate the model object (that creates step_model and train_model)
    model = Model(policy=policy, env=env, nsteps=nsteps, ent_coef=ent_coef, vf_coef=vf_coef,
        max_grad_norm=max_grad_norm, lr=lr, alpha=alpha, epsilon=epsilon, total_timesteps=total_timesteps, lrschedule=lrschedule)
    if load_path is not None:
        model.load(load_path)

    # Instantiate the runner object
    runner = Runner(env, model, nsteps=nsteps, gamma=gamma)
    epinfobuf = deque(maxlen=100)

    # Calculate the batch_size
    nbatch = nenvs*nsteps

    # Start total timer
    tstart = time.time()

    for update in range(1, total_timesteps//nbatch+1):
        # Get mini batch of experiences
        obs, states, rewards, masks, actions, values, epinfos = runner.run()
        epinfobuf.extend(epinfos)

        policy_loss, value_loss, policy_entropy = model.train(obs, states, rewards, masks, actions, values)
        nseconds = time.time()-tstart

        # Calculate the fps (frame per second)
        fps = int((update*nbatch)/nseconds)
        if update % log_interval == 0 or update == 1:
            # Calculates if value function is a good predicator of the returns (ev > 1)
            # or if it's just worse than predicting nothing (ev =< 0)
            ev = explained_variance(values, rewards)
            logger.record_tabular("nupdates", update)
            logger.record_tabular("total_timesteps", update*nbatch)
            logger.record_tabular("fps", fps)
            logger.record_tabular("policy_entropy", float(policy_entropy))
            logger.record_tabular("value_loss", float(value_loss))
            logger.record_tabular("explained_variance", float(ev))
            logger.record_tabular("eprewmean", safemean([epinfo['r'] for epinfo in epinfobuf]))
            logger.record_tabular("eplenmean", safemean([epinfo['l'] for epinfo in epinfobuf]))
            logger.dump_tabular()
    return model

def learn ( network , env , seed = None , nsteps = 5 , total_timesteps = int ( 80e6 ) , vf_coef = 0.5 , ent_coef = 0.01 , max_grad_norm = 0.5 , lr = 7e-4 , lrschedule = 'linear' , epsilon = 1e-5 , alpha = 0.99 , gamma = 0.99 , log_interval = 100 , load_path = None , ** network_kwargs ) : set_global_seeds ( seed ) nenvs = env . num_envs policy = build_policy ( env , network , ** network_kwargs ) model = Model ( policy = policy , env = env , nsteps = nsteps , ent_coef = ent_coef , vf_coef = vf_coef , max_grad_norm = max_grad_norm , lr = lr , alpha = alpha , epsilon = epsilon , total_timesteps = total_timesteps , lrschedule = lrschedule ) if load_path is not None : model . load ( load_path ) runner = Runner ( env , model , nsteps = nsteps , gamma = gamma ) epinfobuf = deque ( maxlen = 100 ) nbatch = nenvs * nsteps tstart = time . time ( ) for update in range ( 1 , total_timesteps // nbatch + 1 ) : obs , states , rewards , masks , actions , values , epinfos = runner . run ( ) epinfobuf . extend ( epinfos ) policy_loss , value_loss , policy_entropy = model . train ( obs , states , rewards , masks , actions , values ) nseconds = time . time ( ) - tstart fps = int ( ( update * nbatch ) / nseconds ) if update % log_interval == 0 or update == 1 : ev = explained_variance ( values , rewards ) logger . record_tabular ( "nupdates" , update ) logger . record_tabular ( "total_timesteps" , update * nbatch ) logger . record_tabular ( "fps" , fps ) logger . record_tabular ( "policy_entropy" , float ( policy_entropy ) ) logger . record_tabular ( "value_loss" , float ( value_loss ) ) logger . record_tabular ( "explained_variance" , float ( ev ) ) logger . record_tabular ( "eprewmean" , safemean ( [ epinfo [ 'r' ] for epinfo in epinfobuf ] ) ) logger . record_tabular ( "eplenmean" , safemean ( [ epinfo [ 'l' ] for epinfo in epinfobuf ] ) ) logger . dump_tabular ( ) return model

Main entrypoint for A2C algorithm . Train a policy with given network architecture on a given environment using a2c algorithm .


def sf01(arr):
    """
    swap and then flatten axes 0 and 1
    """
    s = arr.shape
    return arr.swapaxes(0, 1).reshape(s[0] * s[1], *s[2:])

def sf01 ( arr ) : s = arr . shape return arr . swapaxes ( 0 , 1 ) . reshape ( s [ 0 ] * s [ 1 ] , * s [ 2 : ] )


def pretty_eta(seconds_left):
    """Print the number of seconds in human readable format.

    Examples:
    2 days
    2 hours and 37 minutes
    less than a minute

    Paramters
    ---------
    seconds_left: int
        Number of seconds to be converted to the ETA
    Returns
    -------
    eta: str
        String representing the pretty ETA.
    """
    minutes_left = seconds_left // 60
    seconds_left %= 60
    hours_left = minutes_left // 60
    minutes_left %= 60
    days_left = hours_left // 24
    hours_left %= 24

    def helper(cnt, name):
        return "{} {}{}".format(str(cnt), name, ('s' if cnt > 1 else ''))

    if days_left > 0:
        msg = helper(days_left, 'day')
        if hours_left > 0:
            msg += ' and ' + helper(hours_left, 'hour')
        return msg
    if hours_left > 0:
        msg = helper(hours_left, 'hour')
        if minutes_left > 0:
            msg += ' and ' + helper(minutes_left, 'minute')
        return msg
    if minutes_left > 0:
        return helper(minutes_left, 'minute')
    return 'less than a minute'

def
pretty_eta
(
seconds_left
)
:
minutes_left
=
seconds_left
//
60
seconds_left
%=
60
hours_left
=
minutes_left
//
60
minutes_left
%=
60
days_left
=
hours_left
//
24
hours_left
%=
24
def
helper
(
cnt
,
name
)
:
return
"{} {}{}"
.
format
(
str
(
cnt
)
,
name
,
(
's'
if
cnt
>
1
else
''
)
)
if
days_left
>
0
:
msg
=
helper
(
days_left
,
'day'
)
if
hours_left
>
0
:
msg
+=
' and '
+
helper
(
hours_left
,
'hour'
)
return
msg
if
hours_left
>
0
:
msg
=
helper
(
hours_left
,
'hour'
)
if
minutes_left
>
0
:
msg
+=
' and '
+
helper
(
minutes_left
,
'minute'
)
return
msg
if
minutes_left
>
0
:
return
helper
(
minutes_left
,
'minute'
)
return
'less than a minute'

Print the number of seconds in human readable format .


def boolean_flag(parser, name, default=False, help=None):
    """Add a boolean flag to argparse parser.

    Parameters
    ----------
    parser: argparse.Parser
        parser to add the flag to
    name: str
        --<name> will enable the flag, while --no-<name> will disable it
    default: bool or None
        default value of the flag
    help: str
        help string for the flag
    """
    dest = name.replace('-', '_')
    parser.add_argument("--" + name, action="store_true", default=default, dest=dest, help=help)
    parser.add_argument("--no-" + name, action="store_false", dest=dest)

def boolean_flag ( parser , name , default = False , help = None ) : dest = name . replace ( '-' , '_' ) parser . add_argument ( "--" + name , action = "store_true" , default = default , dest = dest , help = help ) parser . add_argument ( "--no-" + name , action = "store_false" , dest = dest )


def get_wrapper_by_name(env, classname):
    """Given an a gym environment possibly wrapped multiple times, returns a wrapper
    of class named classname or raises ValueError if no such wrapper was applied

    Parameters
    ----------
    env: gym.Env of gym.Wrapper
        gym environment
    classname: str
        name of the wrapper

    Returns
    -------
    wrapper: gym.Wrapper
        wrapper named classname
    """
    currentenv = env
    while True:
        if classname == currentenv.class_name():
            return currentenv
        elif isinstance(currentenv, gym.Wrapper):
            currentenv = currentenv.env
        else:
            raise ValueError("Couldn't find wrapper named %s" % classname)

def
get_wrapper_by_name
(
env
,
classname
)
:
currentenv
=
env
while
True
:
if
classname
==
currentenv
.
class_name
(
)
:
return
currentenv
elif
isinstance
(
currentenv
,
gym
.
Wrapper
)
:
currentenv
=
currentenv
.
env
else
:
raise
ValueError
(
"Couldn't find wrapper named %s"
%
classname
)

Given an a gym environment possibly wrapped multiple times returns a wrapper of class named classname or raises ValueError if no such wrapper was applied


def pickle_load(path, compression=False):
    """Unpickle a possible compressed pickle.

    Parameters
    ----------
    path: str
        path to the output file
    compression: bool
        if true assumes that pickle was compressed when created and attempts decompression.

    Returns
    -------
    obj: object
        the unpickled object
    """

    if compression:
        with zipfile.ZipFile(path, "r", compression=zipfile.ZIP_DEFLATED) as myzip:
            with myzip.open("data") as f:
                return pickle.load(f)
    else:
        with open(path, "rb") as f:
            return pickle.load(f)

def pickle_load ( path , compression = False ) : if compression : with zipfile . ZipFile ( path , "r" , compression = zipfile . ZIP_DEFLATED ) as myzip : with myzip . open ( "data" ) as f : return pickle . load ( f ) else : with open ( path , "rb" ) as f : return pickle . load ( f )


def update(self, new_val):
        """Update the estimate.

        Parameters
        ----------
        new_val: float
            new observated value of estimated quantity.
        """
        if self._value is None:
            self._value = new_val
        else:
            self._value = self._gamma * self._value + (1.0 - self._gamma) * new_val

def update ( self , new_val ) : if self . _value is None : self . _value = new_val else : self . _value = self . _gamma * self . _value + ( 1.0 - self . _gamma ) * new_val


def store_args(method):
    """Stores provided method args as instance attributes.
    """
    argspec = inspect.getfullargspec(method)
    defaults = {}
    if argspec.defaults is not None:
        defaults = dict(
            zip(argspec.args[-len(argspec.defaults):], argspec.defaults))
    if argspec.kwonlydefaults is not None:
        defaults.update(argspec.kwonlydefaults)
    arg_names = argspec.args[1:]

    @functools.wraps(method)
    def wrapper(*positional_args, **keyword_args):
        self = positional_args[0]
        # Get default arg values
        args = defaults.copy()
        # Add provided arg values
        for name, value in zip(arg_names, positional_args[1:]):
            args[name] = value
        args.update(keyword_args)
        self.__dict__.update(args)
        return method(*positional_args, **keyword_args)

    return wrapper

def store_args ( method ) : argspec = inspect . getfullargspec ( method ) defaults = { } if argspec . defaults is not None : defaults = dict ( zip ( argspec . args [ - len ( argspec . defaults ) : ] , argspec . defaults ) ) if argspec . kwonlydefaults is not None : defaults . update ( argspec . kwonlydefaults ) arg_names = argspec . args [ 1 : ] @ functools . wraps ( method ) def wrapper ( * positional_args , ** keyword_args ) : self = positional_args [ 0 ] args = defaults . copy ( ) for name , value in zip ( arg_names , positional_args [ 1 : ] ) : args [ name ] = value args . update ( keyword_args ) self . __dict__ . update ( args ) return method ( * positional_args , ** keyword_args ) return wrapper

Stores provided method args as instance attributes .


def flatten_grads(var_list, grads):
    """Flattens a variables and their gradients.
    """
    return tf.concat([tf.reshape(grad, [U.numel(v)])
                      for (v, grad) in zip(var_list, grads)], 0)

def flatten_grads ( var_list , grads ) : return tf . concat ( [ tf . reshape ( grad , [ U . numel ( v ) ] ) for ( v , grad ) in zip ( var_list , grads ) ] , 0 )

Flattens a variables and their gradients .


def nn(input, layers_sizes, reuse=None, flatten=False, name=""):
    """Creates a simple neural network
    """
    for i, size in enumerate(layers_sizes):
        activation = tf.nn.relu if i < len(layers_sizes) - 1 else None
        input = tf.layers.dense(inputs=input,
                                units=size,
                                kernel_initializer=tf.contrib.layers.xavier_initializer(),
                                reuse=reuse,
                                name=name + '_' + str(i))
        if activation:
            input = activation(input)
    if flatten:
        assert layers_sizes[-1] == 1
        input = tf.reshape(input, [-1])
    return input

def nn ( input , layers_sizes , reuse = None , flatten = False , name = "" ) : for i , size in enumerate ( layers_sizes ) : activation = tf . nn . relu if i < len ( layers_sizes ) - 1 else None input = tf . layers . dense ( inputs = input , units = size , kernel_initializer = tf . contrib . layers . xavier_initializer ( ) , reuse = reuse , name = name + '_' + str ( i ) ) if activation : input = activation ( input ) if flatten : assert layers_sizes [ - 1 ] == 1 input = tf . reshape ( input , [ - 1 ] ) return input


def mpi_fork(n, extra_mpi_args=[]):
    """Re-launches the current script with workers
    Returns "parent" for original parent, "child" for MPI children
    """
    if n <= 1:
        return "child"
    if os.getenv("IN_MPI") is None:
        env = os.environ.copy()
        env.update(
            MKL_NUM_THREADS="1",
            OMP_NUM_THREADS="1",
            IN_MPI="1"
        )
        # "-bind-to core" is crucial for good performance
        args = ["mpirun", "-np", str(n)] + \
            extra_mpi_args + \
            [sys.executable]

        args += sys.argv
        subprocess.check_call(args, env=env)
        return "parent"
    else:
        install_mpi_excepthook()
        return "child"

def mpi_fork ( n , extra_mpi_args = [ ] ) : if n <= 1 : return "child" if os . getenv ( "IN_MPI" ) is None : env = os . environ . copy ( ) env . update ( MKL_NUM_THREADS = "1" , OMP_NUM_THREADS = "1" , IN_MPI = "1" ) args = [ "mpirun" , "-np" , str ( n ) ] + extra_mpi_args + [ sys . executable ] args += sys . argv subprocess . check_call ( args , env = env ) return "parent" else : install_mpi_excepthook ( ) return "child"

Re - launches the current script with workers Returns parent for original parent child for MPI children


def get_session(config=None):
    """Get default session or create one with a given config"""
    sess = tf.get_default_session()
    if sess is None:
        sess = make_session(config=config, make_default=True)
    return sess

def get_session ( config = None ) : sess = tf . get_default_session ( ) if sess is None : sess = make_session ( config = config , make_default = True ) return sess


def initialize():
    """Initialize all the uninitialized variables in the global scope."""
    new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
    get_session().run(tf.variables_initializer(new_variables))
    ALREADY_INITIALIZED.update(new_variables)

def initialize ( ) : new_variables = set ( tf . global_variables ( ) ) - ALREADY_INITIALIZED get_session ( ) . run ( tf . variables_initializer ( new_variables ) ) ALREADY_INITIALIZED . update ( new_variables )

Initialize all the uninitialized variables in the global scope .


def adjust_shape(placeholder, data):
    '''
    adjust shape of the data to the shape of the placeholder if possible.
    If shape is incompatible, AssertionError is thrown

    Parameters:
        placeholder     tensorflow input placeholder

        data            input data to be (potentially) reshaped to be fed into placeholder

    Returns:
        reshaped data
    '''

    if not isinstance(data, np.ndarray) and not isinstance(data, list):
        return data
    if isinstance(data, list):
        data = np.array(data)

    placeholder_shape = [x or -1 for x in placeholder.shape.as_list()]

    assert _check_shape(placeholder_shape, data.shape), \
        'Shape of data {} is not compatible with shape of the placeholder {}'.format(data.shape, placeholder_shape)

    return np.reshape(data, placeholder_shape)

adjust shape of the data to the shape of the placeholder if possible . If shape is incompatible AssertionError is thrown


def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):
    """Configure environment for DeepMind-style Atari.
    """
    if episode_life:
        env = EpisodicLifeEnv(env)
    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = WarpFrame(env)
    if scale:
        env = ScaledFloatFrame(env)
    if clip_rewards:
        env = ClipRewardEnv(env)
    if frame_stack:
        env = FrameStack(env, 4)
    return env

def wrap_deepmind ( env , episode_life = True , clip_rewards = True , frame_stack = False , scale = False ) : if episode_life : env = EpisodicLifeEnv ( env ) if 'FIRE' in env . unwrapped . get_action_meanings ( ) : env = FireResetEnv ( env ) env = WarpFrame ( env ) if scale : env = ScaledFloatFrame ( env ) if clip_rewards : env = ClipRewardEnv ( env ) if frame_stack : env = FrameStack ( env , 4 ) return env

Configure environment for DeepMind - style Atari .


def reset(self, **kwargs):
        """Reset only when lives are exhausted.
        This way all states are still reachable even though lives are episodic,
        and the learner need not know about any of this behind-the-scenes.
        """
        if self.was_real_done:
            obs = self.env.reset(**kwargs)
        else:
            # no-op step to advance from terminal/lost life state
            obs, _, _, _ = self.env.step(0)
        self.lives = self.env.unwrapped.ale.lives()
        return obs

def reset ( self , ** kwargs ) : if self . was_real_done : obs = self . env . reset ( ** kwargs ) else : obs , _ , _ , _ = self . env . step ( 0 ) self . lives = self . env . unwrapped . ale . lives ( ) return obs

Reset only when lives are exhausted . This way all states are still reachable even though lives are episodic and the learner need not know about any of this behind - the - scenes .


def gpu_count():
    """
    Count the GPUs on this machine.
    """
    if shutil.which('nvidia-smi') is None:
        return 0
    output = subprocess.check_output(['nvidia-smi', '--query-gpu=gpu_name', '--format=csv'])
    return max(0, len(output.split(b'\n')) - 2)

def gpu_count ( ) : if shutil . which ( 'nvidia-smi' ) is None : return 0 output = subprocess . check_output ( [ 'nvidia-smi' , '--query-gpu=gpu_name' , '--format=csv' ] ) return max ( 0 , len ( output . split ( b'\n' ) ) - 2 )


def setup_mpi_gpus():
    """
    Set CUDA_VISIBLE_DEVICES to MPI rank if not already set
    """
    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
        if sys.platform == 'darwin': # This Assumes if you're on OSX you're just
            ids = []                 # doing a smoke test and don't want GPUs
        else:
            lrank, _lsize = get_local_rank_size(MPI.COMM_WORLD)
            ids = [lrank]
        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, ids))

def setup_mpi_gpus ( ) : if 'CUDA_VISIBLE_DEVICES' not in os . environ : if sys . platform == 'darwin' : ids = [ ] else : lrank , _lsize = get_local_rank_size ( MPI . COMM_WORLD ) ids = [ lrank ] os . environ [ "CUDA_VISIBLE_DEVICES" ] = "," . join ( map ( str , ids ) )


def get_local_rank_size(comm):
    """
    Returns the rank of each process on its machine
    The processes on a given machine will be assigned ranks
        0, 1, 2, ..., N-1,
    where N is the number of processes on this machine.

    Useful if you want to assign one gpu per machine
    """
    this_node = platform.node()
    ranks_nodes = comm.allgather((comm.Get_rank(), this_node))
    node2rankssofar = defaultdict(int)
    local_rank = None
    for (rank, node) in ranks_nodes:
        if rank == comm.Get_rank():
            local_rank = node2rankssofar[node]
        node2rankssofar[node] += 1
    assert local_rank is not None
    return local_rank, node2rankssofar[this_node]

def get_local_rank_size ( comm ) : this_node = platform . node ( ) ranks_nodes = comm . allgather ( ( comm . Get_rank ( ) , this_node ) ) node2rankssofar = defaultdict ( int ) local_rank = None for ( rank , node ) in ranks_nodes : if rank == comm . Get_rank ( ) : local_rank = node2rankssofar [ node ] node2rankssofar [ node ] += 1 assert local_rank is not None return local_rank , node2rankssofar [ this_node ]

Returns the rank of each process on its machine The processes on a given machine will be assigned ranks 0 1 2 ... N - 1 where N is the number of processes on this machine .


def share_file(comm, path):
    """
    Copies the file from rank 0 to all other ranks
    Puts it in the same place on all machines
    """
    localrank, _ = get_local_rank_size(comm)
    if comm.Get_rank() == 0:
        with open(path, 'rb') as fh:
            data = fh.read()
        comm.bcast(data)
    else:
        data = comm.bcast(None)
        if localrank == 0:
            os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(path, 'wb') as fh:
                fh.write(data)
    comm.Barrier()

def share_file ( comm , path ) : localrank , _ = get_local_rank_size ( comm ) if comm . Get_rank ( ) == 0 : with open ( path , 'rb' ) as fh : data = fh . read ( ) comm . bcast ( data ) else : data = comm . bcast ( None ) if localrank == 0 : os . makedirs ( os . path . dirname ( path ) , exist_ok = True ) with open ( path , 'wb' ) as fh : fh . write ( data ) comm . Barrier ( )


def dict_gather(comm, d, op='mean', assert_all_have_data=True):
    """
    Perform a reduction operation over dicts
    """
    if comm is None: return d
    alldicts = comm.allgather(d)
    size = comm.size
    k2li = defaultdict(list)
    for d in alldicts:
        for (k,v) in d.items():
            k2li[k].append(v)
    result = {}
    for (k,li) in k2li.items():
        if assert_all_have_data:
            assert len(li)==size, "only %i out of %i MPI workers have sent '%s'" % (len(li), size, k)
        if op=='mean':
            result[k] = np.mean(li, axis=0)
        elif op=='sum':
            result[k] = np.sum(li, axis=0)
        else:
            assert 0, op
    return result


def discount(x, gamma):
    """
    computes discounted sums along 0th dimension of x.

    inputs
    ------
    x: ndarray
    gamma: float

    outputs
    -------
    y: ndarray with same shape as x, satisfying

        y[t] = x[t] + gamma*x[t+1] + gamma^2*x[t+2] + ... + gamma^k x[t+k],
                where k = len(x) - t - 1

    """
    assert x.ndim >= 1
    return scipy.signal.lfilter([1],[1,-gamma],x[::-1], axis=0)[::-1]

def discount ( x , gamma ) : assert x . ndim >= 1 return scipy . signal . lfilter ( [ 1 ] , [ 1 , - gamma ] , x [ : : - 1 ] , axis = 0 ) [ : : - 1 ]

computes discounted sums along 0th dimension of x .


def add(self, *args, **kwargs):
        """See ReplayBuffer.store_effect"""
        idx = self._next_idx
        super().add(*args, **kwargs)
        self._it_sum[idx] = self._max_priority ** self._alpha
        self._it_min[idx] = self._max_priority ** self._alpha

def add ( self , * args , ** kwargs ) : idx = self . _next_idx super ( ) . add ( * args , ** kwargs ) self . _it_sum [ idx ] = self . _max_priority ** self . _alpha self . _it_min [ idx ] = self . _max_priority ** self . _alpha


def update_priorities(self, idxes, priorities):
        """Update priorities of sampled transitions.

        sets priority of transition at index idxes[i] in buffer
        to priorities[i].

        Parameters
        ----------
        idxes: [int]
            List of idxes of sampled transitions
        priorities: [float]
            List of updated priorities corresponding to
            transitions at the sampled idxes denoted by
            variable `idxes`.
        """
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            assert priority > 0
            assert 0 <= idx < len(self._storage)
            self._it_sum[idx] = priority ** self._alpha
            self._it_min[idx] = priority ** self._alpha

            self._max_priority = max(self._max_priority, priority)

def update_priorities ( self , idxes , priorities ) : assert len ( idxes ) == len ( priorities ) for idx , priority in zip ( idxes , priorities ) : assert priority > 0 assert 0 <= idx < len ( self . _storage ) self . _it_sum [ idx ] = priority ** self . _alpha self . _it_min [ idx ] = priority ** self . _alpha self . _max_priority = max ( self . _max_priority , priority )

Update priorities of sampled transitions .


def wrap_deepmind_retro(env, scale=True, frame_stack=4):
    """
    Configure environment for retro games, using config similar to DeepMind-style Atari in wrap_deepmind
    """
    env = WarpFrame(env)
    env = ClipRewardEnv(env)
    if frame_stack > 1:
        env = FrameStack(env, frame_stack)
    if scale:
        env = ScaledFloatFrame(env)
    return env

def wrap_deepmind_retro ( env , scale = True , frame_stack = 4 ) : env = WarpFrame ( env ) env = ClipRewardEnv ( env ) if frame_stack > 1 : env = FrameStack ( env , frame_stack ) if scale : env = ScaledFloatFrame ( env ) return env

Configure environment for retro games using config similar to DeepMind - style Atari in wrap_deepmind


def make_sample_her_transitions(replay_strategy, replay_k, reward_fun):
    """Creates a sample function that can be used for HER experience replay.

    Args:
        replay_strategy (in ['future', 'none']): the HER replay strategy; if set to 'none',
            regular DDPG experience replay is used
        replay_k (int): the ratio between HER replays and regular replays (e.g. k = 4 -> 4 times
            as many HER replays as regular replays are used)
        reward_fun (function): function to re-compute the reward with substituted goals
    """
    if replay_strategy == 'future':
        future_p = 1 - (1. / (1 + replay_k))
    else:  # 'replay_strategy' == 'none'
        future_p = 0

    def _sample_her_transitions(episode_batch, batch_size_in_transitions):
        """episode_batch is {key: array(buffer_size x T x dim_key)}
        """
        T = episode_batch['u'].shape[1]
        rollout_batch_size = episode_batch['u'].shape[0]
        batch_size = batch_size_in_transitions

        # Select which episodes and time steps to use.
        episode_idxs = np.random.randint(0, rollout_batch_size, batch_size)
        t_samples = np.random.randint(T, size=batch_size)
        transitions = {key: episode_batch[key][episode_idxs, t_samples].copy()
                       for key in episode_batch.keys()}

        # Select future time indexes proportional with probability future_p. These
        # will be used for HER replay by substituting in future goals.
        her_indexes = np.where(np.random.uniform(size=batch_size) < future_p)
        future_offset = np.random.uniform(size=batch_size) * (T - t_samples)
        future_offset = future_offset.astype(int)
        future_t = (t_samples + 1 + future_offset)[her_indexes]

        # Replace goal with achieved goal but only for the previously-selected
        # HER transitions (as defined by her_indexes). For the other transitions,
        # keep the original goal.
        future_ag = episode_batch['ag'][episode_idxs[her_indexes], future_t]
        transitions['g'][her_indexes] = future_ag

        # Reconstruct info dictionary for reward  computation.
        info = {}
        for key, value in transitions.items():
            if key.startswith('info_'):
                info[key.replace('info_', '')] = value

        # Re-compute reward since we may have substituted the goal.
        reward_params = {k: transitions[k] for k in ['ag_2', 'g']}
        reward_params['info'] = info
        transitions['r'] = reward_fun(**reward_params)

        transitions = {k: transitions[k].reshape(batch_size, *transitions[k].shape[1:])
                       for k in transitions.keys()}

        assert(transitions['u'].shape[0] == batch_size_in_transitions)

        return transitions

    return _sample_her_transitions

def make_sample_her_transitions ( replay_strategy , replay_k , reward_fun ) : if replay_strategy == 'future' : future_p = 1 - ( 1. / ( 1 + replay_k ) ) else : future_p = 0 def _sample_her_transitions ( episode_batch , batch_size_in_transitions ) : T = episode_batch [ 'u' ] . shape [ 1 ] rollout_batch_size = episode_batch [ 'u' ] . shape [ 0 ] batch_size = batch_size_in_transitions episode_idxs = np . random . randint ( 0 , rollout_batch_size , batch_size ) t_samples = np . random . randint ( T , size = batch_size ) transitions = { key : episode_batch [ key ] [ episode_idxs , t_samples ] . copy ( ) for key in episode_batch . keys ( ) } her_indexes = np . where ( np . random . uniform ( size = batch_size ) < future_p ) future_offset = np . random . uniform ( size = batch_size ) * ( T - t_samples ) future_offset = future_offset . astype ( int ) future_t = ( t_samples + 1 + future_offset ) [ her_indexes ] future_ag = episode_batch [ 'ag' ] [ episode_idxs [ her_indexes ] , future_t ] transitions [ 'g' ] [ her_indexes ] = future_ag info = { } for key , value in transitions . items ( ) : if key . startswith ( 'info_' ) : info [ key . replace ( 'info_' , '' ) ] = value reward_params = { k : transitions [ k ] for k in [ 'ag_2' , 'g' ] } reward_params [ 'info' ] = info transitions [ 'r' ] = reward_fun ( ** reward_params ) transitions = { k : transitions [ k ] . reshape ( batch_size , * transitions [ k ] . shape [ 1 : ] ) for k in transitions . keys ( ) } assert ( transitions [ 'u' ] . shape [ 0 ] == batch_size_in_transitions ) return transitions return _sample_her_transitions

Creates a sample function that can be used for HER experience replay .


def parse_cmdline_kwargs(args):
    '''
    convert a list of '='-spaced command-line arguments to a dictionary, evaluating python objects when possible
    '''
    def parse(v):

        assert isinstance(v, str)
        try:
            return eval(v)
        except (NameError, SyntaxError):
            return v

    return {k: parse(v) for k,v in parse_unknown_args(args).items()}

def parse_cmdline_kwargs ( args ) : def parse ( v ) : assert isinstance ( v , str ) try : return eval ( v ) except ( NameError , SyntaxError ) : return v return { k : parse ( v ) for k , v in parse_unknown_args ( args ) . items ( ) }

convert a list of = - spaced command - line arguments to a dictionary evaluating python objects when possible


def compute_geometric_median(X, eps=1e-5):
    """
    Estimate the geometric median of points in 2D.

    Code from https://stackoverflow.com/a/30305181

    Parameters
    ----------
    X : (N,2) ndarray
        Points in 2D. Second axis must be given in xy-form.

    eps : float, optional
        Distance threshold when to return the median.

    Returns
    -------
    (2,) ndarray
        Geometric median as xy-coordinate.

    """
    y = np.mean(X, 0)

    while True:
        D = scipy.spatial.distance.cdist(X, [y])
        nonzeros = (D != 0)[:, 0]

        Dinv = 1 / D[nonzeros]
        Dinvs = np.sum(Dinv)
        W = Dinv / Dinvs
        T = np.sum(W * X[nonzeros], 0)

        num_zeros = len(X) - np.sum(nonzeros)
        if num_zeros == 0:
            y1 = T
        elif num_zeros == len(X):
            return y
        else:
            R = (T - y) * Dinvs
            r = np.linalg.norm(R)
            rinv = 0 if r == 0 else num_zeros/r
            y1 = max(0, 1-rinv)*T + min(1, rinv)*y

        if scipy.spatial.distance.euclidean(y, y1) < eps:
            return y1

        y = y1

def compute_geometric_median ( X , eps = 1e-5 ) : y = np . mean ( X , 0 ) while True : D = scipy . spatial . distance . cdist ( X , [ y ] ) nonzeros = ( D != 0 ) [ : , 0 ] Dinv = 1 / D [ nonzeros ] Dinvs = np . sum ( Dinv ) W = Dinv / Dinvs T = np . sum ( W * X [ nonzeros ] , 0 ) num_zeros = len ( X ) - np . sum ( nonzeros ) if num_zeros == 0 : y1 = T elif num_zeros == len ( X ) : return y else : R = ( T - y ) * Dinvs r = np . linalg . norm ( R ) rinv = 0 if r == 0 else num_zeros / r y1 = max ( 0 , 1 - rinv ) * T + min ( 1 , rinv ) * y if scipy . spatial . distance . euclidean ( y , y1 ) < eps : return y1 y = y1

Estimate the geometric median of points in 2D .


def project(self, from_shape, to_shape):
        """
        Project the keypoint onto a new position on a new image.

        E.g. if the keypoint is on its original image at x=(10 of 100 pixels)
        and y=(20 of 100 pixels) and is projected onto a new image with
        size (width=200, height=200), its new position will be (20, 40).

        This is intended for cases where the original image is resized.
        It cannot be used for more complex changes (e.g. padding, cropping).

        Parameters
        ----------
        from_shape : tuple of int
            Shape of the original image. (Before resize.)

        to_shape : tuple of int
            Shape of the new image. (After resize.)

        Returns
        -------
        imgaug.Keypoint
            Keypoint object with new coordinates.

        """
        xy_proj = project_coords([(self.x, self.y)], from_shape, to_shape)
        return self.deepcopy(x=xy_proj[0][0], y=xy_proj[0][1])

def project ( self , from_shape , to_shape ) : xy_proj = project_coords ( [ ( self . x , self . y ) ] , from_shape , to_shape ) return self . deepcopy ( x = xy_proj [ 0 ] [ 0 ] , y = xy_proj [ 0 ] [ 1 ] )

Project the keypoint onto a new position on a new image .


def shift(self, x=0, y=0):
        """
        Move the keypoint around on an image.

        Parameters
        ----------
        x : number, optional
            Move by this value on the x axis.

        y : number, optional
            Move by this value on the y axis.

        Returns
        -------
        imgaug.Keypoint
            Keypoint object with new coordinates.

        """
        return self.deepcopy(self.x + x, self.y + y)

def shift ( self , x = 0 , y = 0 ) : return self . deepcopy ( self . x + x , self . y + y )


def draw_on_image(self, image, color=(0, 255, 0), alpha=1.0, size=3,
                      copy=True, raise_if_out_of_image=False):
        """
        Draw the keypoint onto a given image.

        The keypoint is drawn as a square.

        Parameters
        ----------
        image : (H,W,3) ndarray
            The image onto which to draw the keypoint.

        color : int or list of int or tuple of int or (3,) ndarray, optional
            The RGB color of the keypoint. If a single int ``C``, then that is
            equivalent to ``(C,C,C)``.

        alpha : float, optional
            The opacity of the drawn keypoint, where ``1.0`` denotes a fully
            visible keypoint and ``0.0`` an invisible one.

        size : int, optional
            The size of the keypoint. If set to ``S``, each square will have
            size ``S x S``.

        copy : bool, optional
            Whether to copy the image before drawing the keypoint.

        raise_if_out_of_image : bool, optional
            Whether to raise an exception if the keypoint is outside of the
            image.

        Returns
        -------
        image : (H,W,3) ndarray
            Image with drawn keypoint.

        """
        if copy:
            image = np.copy(image)

        if image.ndim == 2:
            assert ia.is_single_number(color), (
                "Got a 2D image. Expected then 'color' to be a single number, "
                "but got %s." % (str(color),))
        elif image.ndim == 3 and ia.is_single_number(color):
            color = [color] * image.shape[-1]

        input_dtype = image.dtype
        alpha_color = color
        if alpha < 0.01:
            # keypoint invisible, nothing to do
            return image
        elif alpha > 0.99:
            alpha = 1
        else:
            image = image.astype(np.float32, copy=False)
            alpha_color = alpha * np.array(color)

        height, width = image.shape[0:2]

        y, x = self.y_int, self.x_int

        x1 = max(x - size//2, 0)
        x2 = min(x + 1 + size//2, width)
        y1 = max(y - size//2, 0)
        y2 = min(y + 1 + size//2, height)

        x1_clipped, x2_clipped = np.clip([x1, x2], 0, width)
        y1_clipped, y2_clipped = np.clip([y1, y2], 0, height)

        x1_clipped_ooi = (x1_clipped < 0 or x1_clipped >= width)
        x2_clipped_ooi = (x2_clipped < 0 or x2_clipped >= width+1)
        y1_clipped_ooi = (y1_clipped < 0 or y1_clipped >= height)
        y2_clipped_ooi = (y2_clipped < 0 or y2_clipped >= height+1)
        x_ooi = (x1_clipped_ooi and x2_clipped_ooi)
        y_ooi = (y1_clipped_ooi and y2_clipped_ooi)
        x_zero_size = (x2_clipped - x1_clipped) < 1  # min size is 1px
        y_zero_size = (y2_clipped - y1_clipped) < 1
        if not x_ooi and not y_ooi and not x_zero_size and not y_zero_size:
            if alpha == 1:
                image[y1_clipped:y2_clipped, x1_clipped:x2_clipped] = color
            else:
                image[y1_clipped:y2_clipped, x1_clipped:x2_clipped] = (
                        (1 - alpha)
                        * image[y1_clipped:y2_clipped, x1_clipped:x2_clipped]
                        + alpha_color
                )
        else:
            if raise_if_out_of_image:
                raise Exception(
                    "Cannot draw keypoint x=%.8f, y=%.8f on image with "
                    "shape %s." % (y, x, image.shape))

        if image.dtype.name != input_dtype.name:
            if input_dtype.name == "uint8":
                image = np.clip(image, 0, 255, out=image)
            image = image.astype(input_dtype, copy=False)
        return image


def generate_similar_points_manhattan(self, nb_steps, step_size, return_array=False):
        """
        Generate nearby points to this keypoint based on manhattan distance.

        To generate the first neighbouring points, a distance of S (step size) is moved from the
        center point (this keypoint) to the top, right, bottom and left, resulting in four new
        points. From these new points, the pattern is repeated. Overlapping points are ignored.

        The resulting points have a shape similar to a square rotated by 45 degrees.

        Parameters
        ----------
        nb_steps : int
            The number of steps to move from the center point. nb_steps=1 results in a total of
            5 output points (1 center point + 4 neighbours).

        step_size : number
            The step size to move from every point to its neighbours.

        return_array : bool, optional
            Whether to return the generated points as a list of keypoints or an array
            of shape ``(N,2)``, where ``N`` is the number of generated points and the second axis contains
            the x- (first value) and y- (second value) coordinates.

        Returns
        -------
        points : list of imgaug.Keypoint or (N,2) ndarray
            If return_array was False, then a list of Keypoint.
            Otherwise a numpy array of shape ``(N,2)``, where ``N`` is the number of generated points and
            the second axis contains the x- (first value) and y- (second value) coordinates.
            The center keypoint (the one on which this function was called) is always included.

        """
        # TODO add test
        # Points generates in manhattan style with S steps have a shape similar to a 45deg rotated
        # square. The center line with the origin point has S+1+S = 1+2*S points (S to the left,
        # S to the right). The lines above contain (S+1+S)-2 + (S+1+S)-2-2 + ... + 1 points. E.g.
        # for S=2 it would be 3+1=4 and for S=3 it would be 5+3+1=9. Same for the lines below the
        # center. Hence the total number of points is S+1+S + 2*(S^2).
        points = np.zeros((nb_steps + 1 + nb_steps + 2*(nb_steps**2), 2), dtype=np.float32)

        # we start at the bottom-most line and move towards the top-most line
        yy = np.linspace(self.y - nb_steps * step_size, self.y + nb_steps * step_size, nb_steps + 1 + nb_steps)

        # bottom-most line contains only one point
        width = 1

        nth_point = 0
        for i_y, y in enumerate(yy):
            if width == 1:
                xx = [self.x]
            else:
                xx = np.linspace(self.x - (width-1)//2 * step_size, self.x + (width-1)//2 * step_size, width)
            for x in xx:
                points[nth_point] = [x, y]
                nth_point += 1
            if i_y < nb_steps:
                width += 2
            else:
                width -= 2

        if return_array:
            return points
        return [self.deepcopy(x=points[i, 0], y=points[i, 1]) for i in sm.xrange(points.shape[0])]

def generate_similar_points_manhattan ( self , nb_steps , step_size , return_array = False ) : points = np . zeros ( ( nb_steps + 1 + nb_steps + 2 * ( nb_steps ** 2 ) , 2 ) , dtype = np . float32 ) yy = np . linspace ( self . y - nb_steps * step_size , self . y + nb_steps * step_size , nb_steps + 1 + nb_steps ) width = 1 nth_point = 0 for i_y , y in enumerate ( yy ) : if width == 1 : xx = [ self . x ] else : xx = np . linspace ( self . x - ( width - 1 ) // 2 * step_size , self . x + ( width - 1 ) // 2 * step_size , width ) for x in xx : points [ nth_point ] = [ x , y ] nth_point += 1 if i_y < nb_steps : width += 2 else : width -= 2 if return_array : return points return [ self . deepcopy ( x = points [ i , 0 ] , y = points [ i , 1 ] ) for i in sm . xrange ( points . shape [ 0 ] ) ]

Generate nearby points to this keypoint based on manhattan distance .


def copy(self, x=None, y=None):
        """
        Create a shallow copy of the Keypoint object.

        Parameters
        ----------
        x : None or number, optional
            Coordinate of the keypoint on the x axis.
            If ``None``, the instance's value will be copied.

        y : None or number, optional
            Coordinate of the keypoint on the y axis.
            If ``None``, the instance's value will be copied.

        Returns
        -------
        imgaug.Keypoint
            Shallow copy.

        """
        return self.deepcopy(x=x, y=y)

def copy ( self , x = None , y = None ) : return self . deepcopy ( x = x , y = y )

Create a shallow copy of the Keypoint object .


def deepcopy(self, x=None, y=None):
        """
        Create a deep copy of the Keypoint object.

        Parameters
        ----------
        x : None or number, optional
            Coordinate of the keypoint on the x axis.
            If ``None``, the instance's value will be copied.

        y : None or number, optional
            Coordinate of the keypoint on the y axis.
            If ``None``, the instance's value will be copied.

        Returns
        -------
        imgaug.Keypoint
            Deep copy.

        """
        x = self.x if x is None else x
        y = self.y if y is None else y
        return Keypoint(x=x, y=y)

def deepcopy ( self , x = None , y = None ) : x = self . x if x is None else x y = self . y if y is None else y return Keypoint ( x = x , y = y )

Create a deep copy of the Keypoint object .


def on(self, image):
        """
        Project keypoints from one image to a new one.

        Parameters
        ----------
        image : ndarray or tuple of int
            New image onto which the keypoints are to be projected.
            May also simply be that new image's shape tuple.

        Returns
        -------
        keypoints : imgaug.KeypointsOnImage
            Object containing all projected keypoints.

        """
        shape = normalize_shape(image)
        if shape[0:2] == self.shape[0:2]:
            return self.deepcopy()
        else:
            keypoints = [kp.project(self.shape, shape) for kp in self.keypoints]
            return self.deepcopy(keypoints, shape)

def on ( self , image ) : shape = normalize_shape ( image ) if shape [ 0 : 2 ] == self . shape [ 0 : 2 ] : return self . deepcopy ( ) else : keypoints = [ kp . project ( self . shape , shape ) for kp in self . keypoints ] return self . deepcopy ( keypoints , shape )

Project keypoints from one image to a new one .


def draw_on_image(self, image, color=(0, 255, 0), alpha=1.0, size=3,
                      copy=True, raise_if_out_of_image=False):
        """
        Draw all keypoints onto a given image.

        Each keypoint is marked by a square of a chosen color and size.

        Parameters
        ----------
        image : (H,W,3) ndarray
            The image onto which to draw the keypoints.
            This image should usually have the same shape as
            set in KeypointsOnImage.shape.

        color : int or list of int or tuple of int or (3,) ndarray, optional
            The RGB color of all keypoints. If a single int ``C``, then that is
            equivalent to ``(C,C,C)``.

        alpha : float, optional
            The opacity of the drawn keypoint, where ``1.0`` denotes a fully
            visible keypoint and ``0.0`` an invisible one.

        size : int, optional
            The size of each point. If set to ``C``, each square will have
            size ``C x C``.

        copy : bool, optional
            Whether to copy the image before drawing the points.

        raise_if_out_of_image : bool, optional
            Whether to raise an exception if any keypoint is outside of the image.

        Returns
        -------
        image : (H,W,3) ndarray
            Image with drawn keypoints.

        """
        image = np.copy(image) if copy else image
        for keypoint in self.keypoints:
            image = keypoint.draw_on_image(
                image, color=color, alpha=alpha, size=size, copy=False,
                raise_if_out_of_image=raise_if_out_of_image)
        return image

def draw_on_image ( self , image , color = ( 0 , 255 , 0 ) , alpha = 1.0 , size = 3 , copy = True , raise_if_out_of_image = False ) : image = np . copy ( image ) if copy else image for keypoint in self . keypoints : image = keypoint . draw_on_image ( image , color = color , alpha = alpha , size = size , copy = False , raise_if_out_of_image = raise_if_out_of_image ) return image


def shift(self, x=0, y=0):
        """
        Move the keypoints around on an image.

        Parameters
        ----------
        x : number, optional
            Move each keypoint by this value on the x axis.

        y : number, optional
            Move each keypoint by this value on the y axis.

        Returns
        -------
        out : KeypointsOnImage
            Keypoints after moving them.

        """
        keypoints = [keypoint.shift(x=x, y=y) for keypoint in self.keypoints]
        return self.deepcopy(keypoints)

def shift ( self , x = 0 , y = 0 ) : keypoints = [ keypoint . shift ( x = x , y = y ) for keypoint in self . keypoints ] return self . deepcopy ( keypoints )


def copy(self, keypoints=None, shape=None):
        """
        Create a shallow copy of the KeypointsOnImage object.

        Parameters
        ----------
        keypoints : None or list of imgaug.Keypoint, optional
            List of keypoints on the image. If ``None``, the instance's
            keypoints will be copied.

        shape : tuple of int, optional
            The shape of the image on which the keypoints are placed.
            If ``None``, the instance's shape will be copied.

        Returns
        -------
        imgaug.KeypointsOnImage
            Shallow copy.

        """
        result = copy.copy(self)
        if keypoints is not None:
            result.keypoints = keypoints
        if shape is not None:
            result.shape = shape
        return result

def copy ( self , keypoints = None , shape = None ) : result = copy . copy ( self ) if keypoints is not None : result . keypoints = keypoints if shape is not None : result . shape = shape return result

Create a shallow copy of the KeypointsOnImage object .


def deepcopy(self, keypoints=None, shape=None):
        """
        Create a deep copy of the KeypointsOnImage object.

        Parameters
        ----------
        keypoints : None or list of imgaug.Keypoint, optional
            List of keypoints on the image. If ``None``, the instance's
            keypoints will be copied.

        shape : tuple of int, optional
            The shape of the image on which the keypoints are placed.
            If ``None``, the instance's shape will be copied.

        Returns
        -------
        imgaug.KeypointsOnImage
            Deep copy.

        """
        # for some reason deepcopy is way slower here than manual copy
        if keypoints is None:
            keypoints = [kp.deepcopy() for kp in self.keypoints]
        if shape is None:
            shape = tuple(self.shape)
        return KeypointsOnImage(keypoints, shape)

def deepcopy ( self , keypoints = None , shape = None ) : if keypoints is None : keypoints = [ kp . deepcopy ( ) for kp in self . keypoints ] if shape is None : shape = tuple ( self . shape ) return KeypointsOnImage ( keypoints , shape )

Create a deep copy of the KeypointsOnImage object .


def project(self, from_shape, to_shape):
        """
        Project the bounding box onto a differently shaped image.

        E.g. if the bounding box is on its original image at
        x1=(10 of 100 pixels) and y1=(20 of 100 pixels) and is projected onto
        a new image with size (width=200, height=200), its new position will
        be (x1=20, y1=40). (Analogous for x2/y2.)

        This is intended for cases where the original image is resized.
        It cannot be used for more complex changes (e.g. padding, cropping).

        Parameters
        ----------
        from_shape : tuple of int or ndarray
            Shape of the original image. (Before resize.)

        to_shape : tuple of int or ndarray
            Shape of the new image. (After resize.)

        Returns
        -------
        out : imgaug.BoundingBox
            BoundingBox object with new coordinates.

        """
        coords_proj = project_coords([(self.x1, self.y1), (self.x2, self.y2)],
                                     from_shape, to_shape)
        return self.copy(
            x1=coords_proj[0][0],
            y1=coords_proj[0][1],
            x2=coords_proj[1][0],
            y2=coords_proj[1][1],
            label=self.label)

def project ( self , from_shape , to_shape ) : coords_proj = project_coords ( [ ( self . x1 , self . y1 ) , ( self . x2 , self . y2 ) ] , from_shape , to_shape ) return self . copy ( x1 = coords_proj [ 0 ] [ 0 ] , y1 = coords_proj [ 0 ] [ 1 ] , x2 = coords_proj [ 1 ] [ 0 ] , y2 = coords_proj [ 1 ] [ 1 ] , label = self . label )

Project the bounding box onto a differently shaped image .


def extend(self, all_sides=0, top=0, right=0, bottom=0, left=0):
        """
        Extend the size of the bounding box along its sides.

        Parameters
        ----------
        all_sides : number, optional
            Value by which to extend the bounding box size along all sides.

        top : number, optional
            Value by which to extend the bounding box size along its top side.

        right : number, optional
            Value by which to extend the bounding box size along its right side.

        bottom : number, optional
            Value by which to extend the bounding box size along its bottom side.

        left : number, optional
            Value by which to extend the bounding box size along its left side.

        Returns
        -------
        imgaug.BoundingBox
            Extended bounding box.

        """
        return BoundingBox(
            x1=self.x1 - all_sides - left,
            x2=self.x2 + all_sides + right,
            y1=self.y1 - all_sides - top,
            y2=self.y2 + all_sides + bottom
        )

def extend ( self , all_sides = 0 , top = 0 , right = 0 , bottom = 0 , left = 0 ) : return BoundingBox ( x1 = self . x1 - all_sides - left , x2 = self . x2 + all_sides + right , y1 = self . y1 - all_sides - top , y2 = self . y2 + all_sides + bottom )

Extend the size of the bounding box along its sides .


def intersection(self, other, default=None):
        """
        Compute the intersection bounding box of this bounding box and another one.

        Note that in extreme cases, the intersection can be a single point, meaning that the intersection bounding box
        will exist, but then also has a height and width of zero.

        Parameters
        ----------
        other : imgaug.BoundingBox
            Other bounding box with which to generate the intersection.

        default : any, optional
            Default value to return if there is no intersection.

        Returns
        -------
        imgaug.BoundingBox or any
            Intersection bounding box of the two bounding boxes if there is an intersection.
            If there is no intersection, the default value will be returned, which can by anything.

        """
        x1_i = max(self.x1, other.x1)
        y1_i = max(self.y1, other.y1)
        x2_i = min(self.x2, other.x2)
        y2_i = min(self.y2, other.y2)
        if x1_i > x2_i or y1_i > y2_i:
            return default
        else:
            return BoundingBox(x1=x1_i, y1=y1_i, x2=x2_i, y2=y2_i)

def intersection ( self , other , default = None ) : x1_i = max ( self . x1 , other . x1 ) y1_i = max ( self . y1 , other . y1 ) x2_i = min ( self . x2 , other . x2 ) y2_i = min ( self . y2 , other . y2 ) if x1_i > x2_i or y1_i > y2_i : return default else : return BoundingBox ( x1 = x1_i , y1 = y1_i , x2 = x2_i , y2 = y2_i )

Compute the intersection bounding box of this bounding box and another one .


def union(self, other):
        """
        Compute the union bounding box of this bounding box and another one.

        This is equivalent to drawing a bounding box around all corners points of both
        bounding boxes.

        Parameters
        ----------
        other : imgaug.BoundingBox
            Other bounding box with which to generate the union.

        Returns
        -------
        imgaug.BoundingBox
            Union bounding box of the two bounding boxes.

        """
        return BoundingBox(
            x1=min(self.x1, other.x1),
            y1=min(self.y1, other.y1),
            x2=max(self.x2, other.x2),
            y2=max(self.y2, other.y2),
        )

def union ( self , other ) : return BoundingBox ( x1 = min ( self . x1 , other . x1 ) , y1 = min ( self . y1 , other . y1 ) , x2 = max ( self . x2 , other . x2 ) , y2 = max ( self . y2 , other . y2 ) , )

Compute the union bounding box of this bounding box and another one .


def iou(self, other):
        """
        Compute the IoU of this bounding box with another one.

        IoU is the intersection over union, defined as::

            ``area(intersection(A, B)) / area(union(A, B))``
            ``= area(intersection(A, B)) / (area(A) + area(B) - area(intersection(A, B)))``

        Parameters
        ----------
        other : imgaug.BoundingBox
            Other bounding box with which to compare.

        Returns
        -------
        float
            IoU between the two bounding boxes.

        """
        inters = self.intersection(other)
        if inters is None:
            return 0.0
        else:
            area_union = self.area + other.area - inters.area
            return inters.area / area_union if area_union > 0 else 0.0

def iou ( self , other ) : inters = self . intersection ( other ) if inters is None : return 0.0 else : area_union = self . area + other . area - inters . area return inters . area / area_union if area_union > 0 else 0.0

Compute the IoU of this bounding box with another one .


def is_fully_within_image(self, image):
        """
        Estimate whether the bounding box is fully inside the image area.

        Parameters
        ----------
        image : (H,W,...) ndarray or tuple of int
            Image dimensions to use.
            If an ndarray, its shape will be used.
            If a tuple, it is assumed to represent the image shape
            and must contain at least two integers.

        Returns
        -------
        bool
            True if the bounding box is fully inside the image area. False otherwise.

        """
        shape = normalize_shape(image)
        height, width = shape[0:2]
        return self.x1 >= 0 and self.x2 < width and self.y1 >= 0 and self.y2 < height

def is_fully_within_image ( self , image ) : shape = normalize_shape ( image ) height , width = shape [ 0 : 2 ] return self . x1 >= 0 and self . x2 < width and self . y1 >= 0 and self . y2 < height

Estimate whether the bounding box is fully inside the image area .


def is_partly_within_image(self, image):
        """
        Estimate whether the bounding box is at least partially inside the image area.

        Parameters
        ----------
        image : (H,W,...) ndarray or tuple of int
            Image dimensions to use.
            If an ndarray, its shape will be used.
            If a tuple, it is assumed to represent the image shape
            and must contain at least two integers.

        Returns
        -------
        bool
            True if the bounding box is at least partially inside the image area. False otherwise.

        """
        shape = normalize_shape(image)
        height, width = shape[0:2]
        eps = np.finfo(np.float32).eps
        img_bb = BoundingBox(x1=0, x2=width-eps, y1=0, y2=height-eps)
        return self.intersection(img_bb) is not None

def is_partly_within_image ( self , image ) : shape = normalize_shape ( image ) height , width = shape [ 0 : 2 ] eps = np . finfo ( np . float32 ) . eps img_bb = BoundingBox ( x1 = 0 , x2 = width - eps , y1 = 0 , y2 = height - eps ) return self . intersection ( img_bb ) is not None

Estimate whether the bounding box is at least partially inside the image area .


def is_out_of_image(self, image, fully=True, partly=False):
        """
        Estimate whether the bounding box is partially or fully outside of the image area.

        Parameters
        ----------
        image : (H,W,...) ndarray or tuple of int
            Image dimensions to use. If an ndarray, its shape will be used. If a tuple, it is
            assumed to represent the image shape and must contain at least two integers.

        fully : bool, optional
            Whether to return True if the bounding box is fully outside fo the image area.

        partly : bool, optional
            Whether to return True if the bounding box is at least partially outside fo the
            image area.

        Returns
        -------
        bool
            True if the bounding box is partially/fully outside of the image area, depending
            on defined parameters. False otherwise.

        """
        if self.is_fully_within_image(image):
            return False
        elif self.is_partly_within_image(image):
            return partly
        else:
            return fully

def is_out_of_image ( self , image , fully = True , partly = False ) : if self . is_fully_within_image ( image ) : return False elif self . is_partly_within_image ( image ) : return partly else : return fully

Estimate whether the bounding box is partially or fully outside of the image area .


def clip_out_of_image(self, image):
        """
        Clip off all parts of the bounding box that are outside of the image.

        Parameters
        ----------
        image : (H,W,...) ndarray or tuple of int
            Image dimensions to use for the clipping of the bounding box.
            If an ndarray, its shape will be used.
            If a tuple, it is assumed to represent the image shape and must contain at least two integers.

        Returns
        -------
        result : imgaug.BoundingBox
            Bounding box, clipped to fall within the image dimensions.

        """
        shape = normalize_shape(image)

        height, width = shape[0:2]
        ia.do_assert(height > 0)
        ia.do_assert(width > 0)

        eps = np.finfo(np.float32).eps
        x1 = np.clip(self.x1, 0, width - eps)
        x2 = np.clip(self.x2, 0, width - eps)
        y1 = np.clip(self.y1, 0, height - eps)
        y2 = np.clip(self.y2, 0, height - eps)

        return self.copy(
            x1=x1,
            y1=y1,
            x2=x2,
            y2=y2,
            label=self.label
        )

def clip_out_of_image ( self , image ) : shape = normalize_shape ( image ) height , width = shape [ 0 : 2 ] ia . do_assert ( height > 0 ) ia . do_assert ( width > 0 ) eps = np . finfo ( np . float32 ) . eps x1 = np . clip ( self . x1 , 0 , width - eps ) x2 = np . clip ( self . x2 , 0 , width - eps ) y1 = np . clip ( self . y1 , 0 , height - eps ) y2 = np . clip ( self . y2 , 0 , height - eps ) return self . copy ( x1 = x1 , y1 = y1 , x2 = x2 , y2 = y2 , label = self . label )

Clip off all parts of the bounding box that are outside of the image .


def draw_on_image(self, image, color=(0, 255, 0), alpha=1.0, size=1,
                      copy=True, raise_if_out_of_image=False, thickness=None):
        """
        Draw the bounding box on an image.

        Parameters
        ----------
        image : (H,W,C) ndarray(uint8)
            The image onto which to draw the bounding box.

        color : iterable of int, optional
            The color to use, corresponding to the channel layout of the image. Usually RGB.

        alpha : float, optional
            The transparency of the drawn bounding box, where 1.0 denotes no transparency and
            0.0 is invisible.

        size : int, optional
            The thickness of the bounding box in pixels. If the value is larger than 1, then
            additional pixels will be added around the bounding box (i.e. extension towards the
            outside).

        copy : bool, optional
            Whether to copy the input image or change it in-place.

        raise_if_out_of_image : bool, optional
            Whether to raise an error if the bounding box is fully outside of the
            image. If set to False, no error will be raised and only the parts inside the image
            will be drawn.

        thickness : None or int, optional
            Deprecated.

        Returns
        -------
        result : (H,W,C) ndarray(uint8)
            Image with bounding box drawn on it.

        """
        if thickness is not None:
            ia.warn_deprecated(
                "Usage of argument 'thickness' in BoundingBox.draw_on_image() "
                "is deprecated. The argument was renamed to 'size'."
            )
            size = thickness

        if raise_if_out_of_image and self.is_out_of_image(image):
            raise Exception("Cannot draw bounding box x1=%.8f, y1=%.8f, x2=%.8f, y2=%.8f on image with shape %s." % (
                self.x1, self.y1, self.x2, self.y2, image.shape))

        result = np.copy(image) if copy else image

        if isinstance(color, (tuple, list)):
            color = np.uint8(color)

        for i in range(size):
            y1, y2, x1, x2 = self.y1_int, self.y2_int, self.x1_int, self.x2_int

            # When y values get into the range (H-0.5, H), the *_int functions round them to H.
            # That is technically sensible, but in the case of drawing means that the border lies
            # just barely outside of the image, making the border disappear, even though the BB
            # is fully inside the image. Here we correct for that because of beauty reasons.
            # Same is the case for x coordinates.
            if self.is_fully_within_image(image):
                y1 = np.clip(y1, 0, image.shape[0]-1)
                y2 = np.clip(y2, 0, image.shape[0]-1)
                x1 = np.clip(x1, 0, image.shape[1]-1)
                x2 = np.clip(x2, 0, image.shape[1]-1)

            y = [y1-i, y1-i, y2+i, y2+i]
            x = [x1-i, x2+i, x2+i, x1-i]
            rr, cc = skimage.draw.polygon_perimeter(y, x, shape=result.shape)
            if alpha >= 0.99:
                result[rr, cc, :] = color
            else:
                if ia.is_float_array(result):
                    result[rr, cc, :] = (1 - alpha) * result[rr, cc, :] + alpha * color
                    result = np.clip(result, 0, 255)
                else:
                    input_dtype = result.dtype
                    result = result.astype(np.float32)
                    result[rr, cc, :] = (1 - alpha) * result[rr, cc, :] + alpha * color
                    result = np.clip(result, 0, 255).astype(input_dtype)

        return result


def extract_from_image(self, image, pad=True, pad_max=None,
                           prevent_zero_size=True):
        """
        Extract the image pixels within the bounding box.

        This function will zero-pad the image if the bounding box is partially/fully outside of
        the image.

        Parameters
        ----------
        image : (H,W) ndarray or (H,W,C) ndarray
            The image from which to extract the pixels within the bounding box.

        pad : bool, optional
            Whether to zero-pad the image if the object is partially/fully
            outside of it.

        pad_max : None or int, optional
            The maximum number of pixels that may be zero-paded on any side,
            i.e. if this has value ``N`` the total maximum of added pixels
            is ``4*N``.
            This option exists to prevent extremely large images as a result of
            single points being moved very far away during augmentation.

        prevent_zero_size : bool, optional
            Whether to prevent height or width of the extracted image from becoming zero.
            If this is set to True and height or width of the bounding box is below 1, the height/width will
            be increased to 1. This can be useful to prevent problems, e.g. with image saving or plotting.
            If it is set to False, images will be returned as ``(H', W')`` or ``(H', W', 3)`` with ``H`` or
            ``W`` potentially being 0.

        Returns
        -------
        image : (H',W') ndarray or (H',W',C) ndarray
            Pixels within the bounding box. Zero-padded if the bounding box is partially/fully
            outside of the image. If prevent_zero_size is activated, it is guarantueed that ``H'>0``
            and ``W'>0``, otherwise only ``H'>=0`` and ``W'>=0``.

        """
        pad_top = 0
        pad_right = 0
        pad_bottom = 0
        pad_left = 0

        height, width = image.shape[0], image.shape[1]
        x1, x2, y1, y2 = self.x1_int, self.x2_int, self.y1_int, self.y2_int

        # When y values get into the range (H-0.5, H), the *_int functions round them to H.
        # That is technically sensible, but in the case of extraction leads to a black border,
        # which is both ugly and unexpected after calling cut_out_of_image(). Here we correct for
        # that because of beauty reasons.
        # Same is the case for x coordinates.
        fully_within = self.is_fully_within_image(image)
        if fully_within:
            y1, y2 = np.clip([y1, y2], 0, height-1)
            x1, x2 = np.clip([x1, x2], 0, width-1)

        # TODO add test
        if prevent_zero_size:
            if abs(x2 - x1) < 1:
                x2 = x1 + 1
            if abs(y2 - y1) < 1:
                y2 = y1 + 1

        if pad:
            # if the bb is outside of the image area, the following pads the image
            # first with black pixels until the bb is inside the image
            # and only then extracts the image area
            # TODO probably more efficient to initialize an array of zeros
            # and copy only the portions of the bb into that array that are
            # natively inside the image area
            if x1 < 0:
                pad_left = abs(x1)
                x2 = x2 + pad_left
                width = width + pad_left
                x1 = 0
            if y1 < 0:
                pad_top = abs(y1)
                y2 = y2 + pad_top
                height = height + pad_top
                y1 = 0
            if x2 >= width:
                pad_right = x2 - width
            if y2 >= height:
                pad_bottom = y2 - height

            paddings = [pad_top, pad_right, pad_bottom, pad_left]
            any_padded = any([val > 0 for val in paddings])
            if any_padded:
                if pad_max is None:
                    pad_max = max(paddings)

                image = ia.pad(
                    image,
                    top=min(pad_top, pad_max),
                    right=min(pad_right, pad_max),
                    bottom=min(pad_bottom, pad_max),
                    left=min(pad_left, pad_max)
                )
            return image[y1:y2, x1:x2]
        else:
            within_image = (
                (0, 0, 0, 0)
                <= (x1, y1, x2, y2)
                < (width, height, width, height)
            )
            out_height, out_width = (y2 - y1), (x2 - x1)
            nonzero_height = (out_height > 0)
            nonzero_width = (out_width > 0)
            if within_image and nonzero_height and nonzero_width:
                return image[y1:y2, x1:x2]
            if prevent_zero_size:
                out_height = 1
                out_width = 1
            else:
                out_height = 0
                out_width = 0
            if image.ndim == 2:
                return np.zeros((out_height, out_width), dtype=image.dtype)
            return np.zeros((out_height, out_width, image.shape[-1]),
                            dtype=image.dtype)

def extract_from_image ( self , image , pad = True , pad_max = None , prevent_zero_size = True ) : pad_top = 0 pad_right = 0 pad_bottom = 0 pad_left = 0 height , width = image . shape [ 0 ] , image . shape [ 1 ] x1 , x2 , y1 , y2 = self . x1_int , self . x2_int , self . y1_int , self . y2_int fully_within = self . is_fully_within_image ( image ) if fully_within : y1 , y2 = np . clip ( [ y1 , y2 ] , 0 , height - 1 ) x1 , x2 = np . clip ( [ x1 , x2 ] , 0 , width - 1 ) if prevent_zero_size : if abs ( x2 - x1 ) < 1 : x2 = x1 + 1 if abs ( y2 - y1 ) < 1 : y2 = y1 + 1 if pad : if x1 < 0 : pad_left = abs ( x1 ) x2 = x2 + pad_left width = width + pad_left x1 = 0 if y1 < 0 : pad_top = abs ( y1 ) y2 = y2 + pad_top height = height + pad_top y1 = 0 if x2 >= width : pad_right = x2 - width if y2 >= height : pad_bottom = y2 - height paddings = [ pad_top , pad_right , pad_bottom , pad_left ] any_padded = any ( [ val > 0 for val in paddings ] ) if any_padded : if pad_max is None : pad_max = max ( paddings ) image = ia . pad ( image , top = min ( pad_top , pad_max ) , right = min ( pad_right , pad_max ) , bottom = min ( pad_bottom , pad_max ) , left = min ( pad_left , pad_max ) ) return image [ y1 : y2 , x1 : x2 ] else : within_image = ( ( 0 , 0 , 0 , 0 ) <= ( x1 , y1 , x2 , y2 ) < ( width , height , width , height ) ) out_height , out_width = ( y2 - y1 ) , ( x2 - x1 ) nonzero_height = ( out_height > 0 ) nonzero_width = ( out_width > 0 ) if within_image and nonzero_height and nonzero_width : return image [ y1 : y2 , x1 : x2 ] if prevent_zero_size : out_height = 1 out_width = 1 else : out_height = 0 out_width = 0 if image . ndim == 2 : return np . zeros ( ( out_height , out_width ) , dtype = image . dtype ) return np . zeros ( ( out_height , out_width , image . shape [ - 1 ] ) , dtype = image . dtype )

Extract the image pixels within the bounding box .


def copy(self, x1=None, y1=None, x2=None, y2=None, label=None):
        """
        Create a shallow copy of the BoundingBox object.

        Parameters
        ----------
        x1 : None or number
            If not None, then the x1 coordinate of the copied object will be set to this value.

        y1 : None or number
            If not None, then the y1 coordinate of the copied object will be set to this value.

        x2 : None or number
            If not None, then the x2 coordinate of the copied object will be set to this value.

        y2 : None or number
            If not None, then the y2 coordinate of the copied object will be set to this value.

        label : None or string
            If not None, then the label of the copied object will be set to this value.

        Returns
        -------
        imgaug.BoundingBox
            Shallow copy.

        """
        return BoundingBox(
            x1=self.x1 if x1 is None else x1,
            x2=self.x2 if x2 is None else x2,
            y1=self.y1 if y1 is None else y1,
            y2=self.y2 if y2 is None else y2,
            label=self.label if label is None else label
        )

def copy ( self , x1 = None , y1 = None , x2 = None , y2 = None , label = None ) : return BoundingBox ( x1 = self . x1 if x1 is None else x1 , x2 = self . x2 if x2 is None else x2 , y1 = self . y1 if y1 is None else y1 , y2 = self . y2 if y2 is None else y2 , label = self . label if label is None else label )

Create a shallow copy of the BoundingBox object .


def draw_on_image(self, image, color=(0, 255, 0), alpha=1.0, size=1,
                      copy=True, raise_if_out_of_image=False, thickness=None):
        """
        Draw all bounding boxes onto a given image.

        Parameters
        ----------
        image : (H,W,3) ndarray
            The image onto which to draw the bounding boxes.
            This image should usually have the same shape as
            set in BoundingBoxesOnImage.shape.

        color : int or list of int or tuple of int or (3,) ndarray, optional
            The RGB color of all bounding boxes. If a single int ``C``, then
            that is equivalent to ``(C,C,C)``.

        alpha : float, optional
            Alpha/transparency of the bounding box.

        size : int, optional
            Thickness in pixels.

        copy : bool, optional
            Whether to copy the image before drawing the bounding boxes.

        raise_if_out_of_image : bool, optional
            Whether to raise an exception if any bounding box is outside of the
            image.

        thickness : None or int, optional
            Deprecated.

        Returns
        -------
        image : (H,W,3) ndarray
            Image with drawn bounding boxes.

        """
        image = np.copy(image) if copy else image

        for bb in self.bounding_boxes:
            image = bb.draw_on_image(
                image,
                color=color,
                alpha=alpha,
                size=size,
                copy=False,
                raise_if_out_of_image=raise_if_out_of_image,
                thickness=thickness
            )

        return image

def draw_on_image ( self , image , color = ( 0 , 255 , 0 ) , alpha = 1.0 , size = 1 , copy = True , raise_if_out_of_image = False , thickness = None ) : image = np . copy ( image ) if copy else image for bb in self . bounding_boxes : image = bb . draw_on_image ( image , color = color , alpha = alpha , size = size , copy = False , raise_if_out_of_image = raise_if_out_of_image , thickness = thickness ) return image

Draw all bounding boxes onto a given image .