text:
#!/usr/bin/env python3

# Copyright 2020 The Johns Hopkins University (author: Jiatong Shi, Hailan Lin)

import torch
import torch.nn as nn
import torch.nn.init as init
from torch.nn import functional as F
from torch.nn import Module
from torch.nn.parameter import Parameter
from torch.nn import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn.init import xavier_normal_
from torch.nn.init import constant_
from torch.nn import Dropout
from torch.nn import Linear
from torch.nn import LayerNorm
import numpy as np
import math
import copy

SCALE_WEIGHT = 0.5 ** 0.5

class GatedConv(nn.Module):
    def __init__(self, input_size, width=3, dropout=0.2, nopad=False):
        super(GatedConv, self).__init__()
        self.conv = nn.Conv2d(in_channels=input_size, out_channels=2 * input_size,
                                     kernel_size=(width, 1), stride=(1, 1),
                                     padding=(width // 2 * (1 - nopad), 0))
        init.xavier_uniform_(self.conv.weight, gain=(4 * (1 - dropout)) ** 0.5)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x_var):
        x_var = self.dropout(x_var)
        x_var = self.conv(x_var)
        out, gate = x_var.split(int(x_var.size(1) / 2), 1)
        out = out * torch.sigmoid(gate)
        return out

class StackedCNN(nn.Module):
    """ Stacked CNN class """

    def __init__(self, num_layers, input_size, cnn_kernel_width=3,
                 dropout=0.2):
        super(StackedCNN, self).__init__()
        self.dropout = dropout
        self.num_layers = num_layers
        self.layers = nn.ModuleList()
        for _ in range(num_layers):
            self.layers.append(
                GatedConv(input_size, cnn_kernel_width, dropout))

    def forward(self, x):
        for conv in self.layers:
            x = x + conv(x)
            x *= SCALE_WEIGHT
        return x

class GLU(nn.Module):
    def __init__(self, num_layers, hidden_size,
                 cnn_kernel_width, dropout, input_size):
        super(GLU, self).__init__()
        self.linear = nn.Linear(input_size, hidden_size)
        self.cnn = StackedCNN(num_layers, hidden_size,
                              cnn_kernel_width, dropout)

    def forward(self, emb):
        emb_reshape = emb.view(emb.size(0) * emb.size(1), -1)
        emb_remap = self.linear(emb_reshape)
        emb_remap = emb_remap.view(emb.size(0), emb.size(1), -1)

        emb_remap = _shape_transform(emb_remap)
        out = self.cnn(emb_remap)

        return out.squeeze(3).contiguous()

class PositionalEncoding(nn.Module):
    """ Positional Encoding.
    Modified from
    https://github.com/pytorch/examples/blob/master/word_language_model/model.py
    Args:
        d_model: the embed dim (required).
        dropout: the dropout value (default=0.1).
        max_len: the max. length of the incoming sequence (default=5000).
    """
    def __init__(self, d_model, dropout=0.1, max_len=5000, device='cuda'):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        pe = pe.to(device)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float()
            * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """Inputs of forward function
        Args:
            x: the sequence fed to the positional encoder model (required).
        Shape:
            x: [sequence length, batch size, embed dim]
            output: [sequence length, batch size, embed dim]
        """
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class CBHG(nn.Module):
    """
    CBHG Module
    """
    def __init__(self, hidden_size, K=16, projection_size = 256, num_gru_layers=2, max_pool_kernel_size=2, is_post=False):
        """
        :param hidden_size: dimension of hidden unit
        :param K: # of convolution banks
        :param projection_size: dimension of projection unit
        :param num_gru_layers: # of layers of GRUcell
        :param max_pool_kernel_size: max pooling kernel size
        :param is_post: whether post processing or not
        """
        super(CBHG, self).__init__()
        self.hidden_size = hidden_size
        self.projection_size = projection_size
        self.convbank_list = nn.ModuleList()
        self.convbank_list.append(nn.Conv1d(in_channels=projection_size,
                                                out_channels=hidden_size,
                                                kernel_size=1,
                                                padding=int(np.floor(1/2))))

        for i in range(2, K+1):
            self.convbank_list.append(nn.Conv1d(in_channels=hidden_size,
                                                out_channels=hidden_size,
                                                kernel_size=i,
                                                padding=int(np.floor(i/2))))

        self.batchnorm_list = nn.ModuleList()
        for i in range(1, K+1):
            self.batchnorm_list.append(nn.BatchNorm1d(hidden_size))

        convbank_outdim = hidden_size * K

        self.conv_projection_1 = nn.Conv1d(in_channels=convbank_outdim,
                                             out_channels=hidden_size,
                                             kernel_size=3,
                                             padding=int(np.floor(3 / 2)))
        self.conv_projection_2 = nn.Conv1d(in_channels=hidden_size,
                                               out_channels=projection_size,
                                               kernel_size=3,
                                               padding=int(np.floor(3 / 2)))
        self.batchnorm_proj_1 = nn.BatchNorm1d(hidden_size)

        self.batchnorm_proj_2 = nn.BatchNorm1d(projection_size)

        self.max_pool = nn.MaxPool1d(max_pool_kernel_size, stride=1, padding=1)
        self.highway = Highwaynet(self.projection_size)
        self.gru = nn.GRU(self.projection_size, self.hidden_size // 2, num_layers=num_gru_layers,
                          batch_first=True,
                          bidirectional=True)

    def _conv_fit_dim(self, x, kernel_size=3):
        if kernel_size % 2 == 0:
            return x[:,:,:-1]
        else:
            return x

    def forward(self, input_):

        input_ = input_.contiguous()
        batch_size = input_.size(0)
        total_length = input_.size(-1)

        convbank_list = list()
        convbank_input = input_

        # Convolution bank filters
        for k, (conv, batchnorm) in enumerate(zip(self.convbank_list, self.batchnorm_list)):
            convbank_input = torch.relu(batchnorm(self._conv_fit_dim(conv(convbank_input), k+1).contiguous()))
            convbank_list.append(convbank_input)

        # Concatenate all features
        conv_cat = torch.cat(convbank_list, dim=1)

        # Max pooling
        conv_cat = self.max_pool(conv_cat)[:,:,:-1]

        # Projection
        conv_projection = torch.relu(self.batchnorm_proj_1(self._conv_fit_dim(self.conv_projection_1(conv_cat))))
        conv_projection = self.batchnorm_proj_2(self._conv_fit_dim(self.conv_projection_2(conv_projection))) + input_

        # Highway networks
        highway = self.highway.forward(conv_projection.transpose(1,2))

        # Bidirectional GRU

        self.gru.flatten_parameters()
        out, _ = self.gru(highway)

        return out

class Highwaynet(nn.Module):
    """
    Highway network
    """
    def __init__(self, num_units, num_layers=4):
        """
        :param num_units: dimension of hidden unit
        :param num_layers: # of highway layers
        """
        super(Highwaynet, self).__init__()
        self.num_units = num_units
        self.num_layers = num_layers
        self.gates = nn.ModuleList()
        self.linears = nn.ModuleList()
        for _ in range(self.num_layers):
            self.linears.append(Linear(num_units, num_units))
            self.gates.append(Linear(num_units, num_units))

    def forward(self, input_):

        out = input_

        # highway gated function
        for fc1, fc2 in zip(self.linears, self.gates):

            h = torch.relu(fc1.forward(out))
            t_ = torch.sigmoid(fc2.forward(out))

            c = 1. - t_
            out = h * t_ + out * c

        return out

class MultiheadAttention(nn.Module):
    """
    Multiheaded Scaled Dot Product Attention
    Implements equation:
    MultiHead(Q, K, V) = Concat(head_1,...,head_h)W^O
        where head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
    Similarly to the above, d_k = d_v = d_model / h
    Inputs
      init:
        nheads : integer # of attention heads
        d_model : model dimensionality
        d_head : dimensionality of a single head
      forward:
        query : [batch size, sequence length, d_model]
        key: [batch size, sequence length, d_model]
        value: [batch size, sequence length, d_model]
      unseen_mask: if True, only attend to previous sequence positions
      src_lengths_mask: if True, mask padding based on src_lengths
    Output
      result : [batch_size, sequence length, d_model]
    """

    def __init__(self, nheads, d_model, device="cuda"):
        "Take in model size and number of heads."
        super(MultiheadAttention, self).__init__()
        assert d_model % nheads == 0
        self.d_head = d_model // nheads
        self.nheads = nheads
        self.device = device
        self.Q_fc = nn.Linear(d_model, d_model, bias=False)
        self.K_fc = nn.Linear(d_model, d_model, bias=False)
        self.V_fc = nn.Linear(d_model, d_model, bias=False)
        self.output_fc = nn.Linear(d_model, d_model, bias=False)
        self.attn = None

    def forward(self, query, key, value, unseen_mask=False, src_lengths=None):
        # 1. Fully-connected layer on q, k, v then
        # 2. Split heads on q, k, v
        # (batch_size, seq_len, d_model) -->
        # (batch_size, nheads, seq_len, d_head)
        query = split_heads(self.Q_fc(query), self.nheads)
        key = split_heads(self.K_fc(key), self.nheads)
        value = split_heads(self.V_fc(value), self.nheads)

        # 4. Scaled dot product attention
        # (batch_size, nheads, seq_len, d_head)
        x, self.attn = scaled_dot_prod_attn(
            query=query,
            key=key,
            value=value,
            unseen_mask=unseen_mask,
            src_lengths=src_lengths,
            device=self.device
        )

        # 5. Combine heads
        x = combine_heads(x)

        # 6. Fully-connected layer for output
        return self.output_fc(x), self.attn

class TransformerEncoderLayer(Module):
    """TransformerEncoderLayer is made up of self-attn and feedforward network.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dim_feedforward: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).
    """

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", device="cuda"):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(nhead, d_model, device)
        # Implementation of Feedforward model
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerEncoderLayer, self).__setstate__(state)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """

        src2, att_weight = self.self_attn(src, src, src, unseen_mask=src_mask,
                              src_lengths=src_key_padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src, att_weight

class TransformerGLULayer(Module):
    """TransformerGLULayer is made up of self-attn and feedforward network.

    Args:
        d_model: the number of expected features in the input (required).
        nhead: the number of heads in the multiheadattention models (required).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).
        glu_kernel: the kernel size of glu block

    """

    def __init__(self, d_model, nhead, dropout=0.1, activation="relu",
        glu_kernel=3, device="cuda"):
        super(TransformerGLULayer, self).__init__()
        self.self_attn = MultiheadAttention(nhead, d_model, device)
        # Implementation of Feedforward model
        self.GLU = GLU(1, d_model, glu_kernel, dropout, d_model)

        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = _get_activation_fn(activation)

    def __setstate__(self, state):
        if 'activation' not in state:
            state['activation'] = F.relu
        super(TransformerGLULayer, self).__setstate__(state)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
        r"""Pass the input through the encoder layer.

        Args:
            src: the sequence to the encoder layer (required).
            src_mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        src = self.norm1(src)
        src2, att_weight = self.self_attn(src, src, src, unseen_mask=src_mask,
                              src_lengths=src_key_padding_mask)

        src = src + self.dropout1(src2)
        src = src * SCALE_WEIGHT
        src = self.norm2(src)
        src2 = self.GLU(src)
        src2 = src2.transpose(1, 2)
        src = src + self.dropout2(src2)
        src = src * SCALE_WEIGHT
        return src, att_weight

class TransformerEncoder(Module):
    r"""TransformerEncoder is a stack of N encoder layers

    Args:
        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
        num_layers: the number of sub-encoder-layers in the encoder (required).
        norm: the layer normalization component (optional).
    """
    __constants__ = ['norm']

    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        assert num_layers > 0
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src, mask=None, src_key_padding_mask=None):
        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> Tensor
        r"""Pass the input through the encoder layers in turn.

        Args:
            src: the sequence to the encoder (required).
            mask: the mask for the src sequence (optional).
            src_key_padding_mask: the mask for the src keys per batch (optional).

        Shape:
            see the docs in Transformer class.
        """
        output = src

        for mod in self.layers:
            output, att_weight = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)

        if self.norm is not None:
            output = self.norm(output)

        return output, att_weight

class Transformer(nn.Module):
    """Transformer encoder based Singing Voice synthesis
    Args:
        hidden_state: the number of expected features in the encoder/decoder inputs.
        nhead: the number of heads in the multiheadattention models.
        num_block: the number of sub-encoder-layers in the encoder.
        fc_dim: the dimension of the feedforward network model.
        dropout: the dropout value.
        pos_enc: True if positional encoding is used.
    """
    def __init__(self, input_dim=128, hidden_state=512, output_dim=128, nhead=4,
        num_block=6, fc_dim=512, dropout=0.1, pos_enc=True):
        super(Transformer, self).__init__()
        self.input_fc = nn.Linear(input_dim, hidden_state)

        self.pos_encoder = PositionalEncoding(hidden_state, dropout)
        # define a single transformer encoder layer

        encoder_layer = TransformerEncoderLayer(hidden_state, nhead, fc_dim, dropout)
        self.input_norm = LayerNorm(hidden_state)
        encoder_norm = LayerNorm(hidden_state)
        self.encoder = TransformerEncoder(encoder_layer,
            num_block, encoder_norm)
        self.postnet = PostNet(hidden_state, output_dim, hidden_state)
        self.hidden_state = hidden_state
        self.pos_enc = pos_enc

        self._reset_parameters()

    def _reset_parameters(self):
        """Initiate parameters in the transformer model."""

        for p in self.parameters():
            if p.dim() > 1:
                xavier_uniform_(p)

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        src = torch.transpose(src, 0, 1)
        embed = self.input_fc(src)
        embed = self.input_norm(embed)
        if self.pos_enc:

            embed, att_weight = self.encoder(embed)
            embed = embed * math.sqrt(self.hidden_state)
            embed = self.pos_encoder(embed)
        memory, att_weight = self.encoder(
            embed,
            mask=src_mask,
            src_key_padding_mask=src_key_padding_mask)
        print(memory.size())
        output = self.postnet(memory)
        output = torch.transpose(output, 0, 1)
        return output, att_weight

class PostNet(nn.Module):
    """
    CBHG Network (mel --> linear)
    """
    def __init__(self, input_channel, output_channel, hidden_state):
        super(PostNet, self).__init__()
        self.pre_projection = nn.Conv1d(input_channel, hidden_state, kernel_size=1, stride=1, padding=0,
            dilation=1, bias=True)
        self.cbhg = CBHG(hidden_state, projection_size=hidden_state)
        self.post_projection = nn.Conv1d(hidden_state, output_channel, kernel_size=1, stride=1, padding=0,
            dilation=1, bias=True)

    def forward(self, x):
        x = x.transpose(1, 2)
        x = self.pre_projection(x)
        x = self.cbhg(x).transpose(1, 2)
        output = self.post_projection(x).transpose(1, 2)

        return output

def create_src_lengths_mask(batch_size, src_lengths, max_srclen=None):
    if max_srclen is None:
        max_srclen = src_lengths.max()

    src_indices = torch.arange(0, max_srclen).unsqueeze(0).type_as(src_lengths)
    src_indices = src_indices.expand(batch_size, max_srclen)
    src_lengths = src_lengths.unsqueeze(dim=1).expand(batch_size, max_srclen)
    # returns [batch_size, max_seq_len]
    return (src_indices < src_lengths).int().detach()

def apply_masks(scores, batch_size, unseen_mask, src_lengths, device):
    seq_len = scores.shape[-1]

    # [1, seq_len, seq_len]
    sequence_mask = torch.ones(seq_len, seq_len).unsqueeze(0).int().to(device)
    if unseen_mask:
        # [1, seq_len, seq_len]
        sequence_mask = (
            torch.tril(torch.ones(seq_len, seq_len), diagonal=0).unsqueeze(0).int()
        )

    if src_lengths is not None:
        # [batch_size, 1, seq_len]
        src_lengths_mask = create_src_lengths_mask(
            batch_size=batch_size, src_lengths=src_lengths, max_srclen=seq_len
        ).unsqueeze(-2).to(device)

        # [batch_size, seq_len, seq_len]
        sequence_mask = sequence_mask & src_lengths_mask

    # [batch_size, 1, seq_len, seq_len]
    sequence_mask = sequence_mask.unsqueeze(1)

    scores = scores.masked_fill(sequence_mask == 0, -np.inf)
    return scores

def scaled_dot_prod_attn(query, key, value, unseen_mask=False, src_lengths=None, device="cuda"):
    """
    Scaled Dot Product Attention
    Implements equation:
    Attention(Q, K, V) = softmax(QK^T/\sqrt{d_k})V
    Inputs:
      query : [batch size, nheads, sequence length, d_k]
      key : [batch size, nheads, sequence length, d_k]
      value : [batch size, nheads, sequence length, d_v]
      unseen_mask: if True, only attend to previous sequence positions
      src_lengths_mask: if True, mask padding based on src_lengths
    Outputs:
      attn: [batch size, sequence length, d_v]
    Note that in this implementation d_q = d_k = d_v = dim
    """
    d_k = query.shape[-1]
    scores = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(d_k)
    if unseen_mask or src_lengths is not None:
        scores = apply_masks(
            scores=scores,
            batch_size=query.shape[0],
            unseen_mask=unseen_mask,
            src_lengths=src_lengths,
            device=device
        )
    p_attn = F.softmax(scores, dim=-1)
    return torch.matmul(p_attn, value), p_attn

def split_heads(X, nheads):
    """
    Split heads:
    1) Split (reshape) last dimension (size d_model) into nheads, d_head
    2) Transpose X from (batch size, sequence length, nheads, d_head) to
        (batch size, nheads, sequence length, d_head)
    Inputs:
      X : [batch size, sequence length, nheads * d_head]
      nheads : integer
    Outputs:
      [batch size,  nheads, sequence length, d_head]
    """
    last_dim = X.shape[-1]
    assert last_dim % nheads == 0
    X_last_dim_split = X.view(list(X.shape[:-1]) + [nheads, last_dim // nheads])
    return X_last_dim_split.transpose(1, 2)

def combine_heads(X):
    """
    Combine heads (the inverse of split heads):
    1) Transpose X from (batch size, nheads, sequence length, d_head) to
        (batch size, sequence length, nheads, d_head)
    2) Combine (reshape) last 2 dimensions (nheads, d_head) into 1 (d_model)
    Inputs:
      X : [batch size * nheads, sequence length, d_head]
      nheads : integer
      d_head : integer
    Outputs:
      [batch_size, seq_len, d_model]
    """
    X = X.transpose(1, 2)
    nheads, d_head = X.shape[-2:]
    return X.contiguous().view(list(X.shape[:-2]) + [nheads * d_head])

def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    else:
        raise RuntimeError("activation should be relu/gelu, not %s." % activation)

def _get_clones(module, N):
    return ModuleList([copy.deepcopy(module) for i in range(N)])

def _shape_transform(x):
    """ Tranform the size of the tensors to fit for conv input. """
    return torch.unsqueeze(torch.transpose(x, 1, 2), 3)


text:
from .fixtures import *

def test_zeros_0(mypytest):
    mypytest('''
    import numpy as np
    a = np.zeros(1)
    b = np.zeros((1,2))

    reveal_type(a)  # Revealed type is 'numpy.ndarray[builtins.float, numpy.OneD]'
    reveal_type(b)  # Revealed type is 'numpy.ndarray[builtins.float, numpy.TwoD]'
    ''')

def test_zeros_1(mypytest):
    mypytest('''
    import numpy as np
    a = np.zeros([1,2])
    b = np.zeros([1 for _ in range(2+2)])
    reveal_type(a)  # Revealed type is 'numpy.ndarray[builtins.float, numpy.TwoD]'
    reveal_type(b)  # Revealed type is 'numpy.ndarray[builtins.float, Any]'
    ''')

def test_zeros_2(mypytest):
    mypytest('''
import numpy as np
a = np.zeros(1, dtype=int)
b = np.zeros(1, dtype='int')
c = np.zeros(1, dtype='i')
reveal_type(a)  # Revealed type is 'numpy.ndarray[builtins.int, numpy.OneD]'
reveal_type(b)  # Revealed type is 'numpy.ndarray[builtins.int, numpy.OneD]'
reveal_type(c)  # Revealed type is 'numpy.ndarray[builtins.int, numpy.OneD]'
''')

ext: py

max_stars_repo_name: rmcgibbo/numpy-mypy

max_issues_repo_name: rmcgibbo/numpy-mypy

max_forks_repo_name: rmcgibbo/numpy-mypy

reason: import numpy


text:
import argparse
import torch
import numpy as np
import random as rnd
import os
import json
from torch.utils.data import DataLoader
import loss_utils
import models
import bilevel_coreset
from cl_streaming import summary
from cl_streaming import datagen
from cl_streaming import training
from cl_streaming import ntk_generator

datasets = ['permmnist', 'splitmnist']
methods = ['uniform', 'coreset',
           'kmeans_features', 'kcenter_features', 'kmeans_grads',
           'kmeans_embedding', 'kcenter_embedding', 'kcenter_grads',
           'entropy', 'hardest', 'frcl', 'icarl', 'grad_matching']

def get_kernel_fn(dataset):
    if dataset == 'permmnist':
        return lambda x, y: ntk_generator.generate_fnn_ntk(x.reshape(-1, 28 * 28), y.reshape(-1, 28 * 28))
    else:
        return lambda x, y: ntk_generator.generate_cnn_ntk(x.reshape(-1, 28, 28, 1), y.reshape(-1, 28, 28, 1))

def continual_learning(args):
    nr_epochs = args.nr_epochs
    beta = args.beta
    dataset = args.dataset
    device = args.device
    method = args.method
    samples_per_task = args.samples_per_task
    buffer_size = args.buffer_size
    num_workers = args.num_workers
    pin_memory = device == 'cuda'
    if dataset == 'permmnist':
        generator = datagen.PermutedMnistGenerator(samples_per_task)
    elif dataset == 'splitmnist':
        generator = datagen.SplitMnistGenerator(samples_per_task)

    tasks = []
    train_loaders = []
    test_loaders = []
    for i in range(generator.max_iter):
        X_train, y_train, X_test, y_test = generator.next_task()
        tasks.append((X_train, y_train, X_test, y_test))
        train_data = datagen.NumpyDataset(X_train, y_train)
        train_loaders.append(
            DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=num_workers,
                       pin_memory=pin_memory))
        test_data = datagen.NumpyDataset(X_test, y_test)
        test_loaders.append(
            DataLoader(test_data, batch_size=args.batch_size, num_workers=num_workers, pin_memory=pin_memory))

    nr_classes = 10
    inner_reg = 1e-3

    if dataset == 'permmnist':
        model = models.FNNet(28 * 28, 100, nr_classes).to(device)
    else:
        model = models.ConvNet(nr_classes).to(device)

    training_op = training.Training(model, device, nr_epochs, beta=beta)
    kernel_fn = get_kernel_fn(dataset)

    bc = bilevel_coreset.BilevelCoreset(outer_loss_fn=loss_utils.cross_entropy,
                                        inner_loss_fn=loss_utils.cross_entropy, out_dim=10, max_outer_it=1,
                                        max_inner_it=200, logging_period=1000)
    rs = np.random.RandomState(args.seed)
    for i in range(generator.max_iter):
        training_op.train(train_loaders[i])
        size_per_task = buffer_size // (i + 1)
        for j in range(i):
            (X, y), w = training_op.buffer[j]
            X, y = X[:size_per_task], y[:size_per_task]
            training_op.buffer[j] = ((X, y), np.ones(len(y)))
        X, y, _, _ = tasks[i]
        if method == 'coreset':
            chosen_inds, _, = bc.build_with_representer_proxy_batch(X, y, size_per_task, kernel_fn, cache_kernel=True,
                                                                    start_size=1, inner_reg=inner_reg)
        else:
            summarizer = summary.Summarizer.factory(method, rs)
            chosen_inds = summarizer.build_summary(X, y, size_per_task, method=method, model=model, device=device)
        X, y = X[chosen_inds], y[chosen_inds]
        assert (X.shape[0] == size_per_task)
        training_op.buffer.append(((X, y), np.ones(len(y))))

    result = []
    for k in range(generator.max_iter):
        result.append(training_op.test(test_loaders[k]))
    filename = '{}_{}_{}_{}_{}.txt'.format(dataset, method, buffer_size, beta, seed)
    if not os.path.exists('cl_results'):
        os.makedirs('cl_results')
    with open('cl_results/' + filename, 'w') as outfile:
        json.dump({'test_acc': np.mean(result), 'acc_per_task': result}, outfile)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Continual Learning')
    parser.add_argument('--seed', type=int, default=0, metavar='seed',
                        help='random seed (default: 0)')
    parser.add_argument('--nr_epochs', default=400, type=int)
    parser.add_argument('--beta', default=1, type=float, help='the buffer penalty')
    parser.add_argument('--dataset', default='splitmnist', choices=datasets)
    parser.add_argument('--method', default='coreset', choices=methods)
    parser.add_argument('--device', default='cuda', choices=['cpu', 'cuda'])
    parser.add_argument('--samples_per_task', default=1000, type=int)
    parser.add_argument('--buffer_size', default=100, type=int)
    parser.add_argument('--batch_size', default=256, type=int)
    parser.add_argument('--num_workers', default=0, type=int)
    args = parser.parse_args()
    print(args)
    seed = args.seed

    torch.set_num_threads(1)
    torch.manual_seed(seed)
    np.random.seed(seed)
    rnd.seed(seed)

    continual_learning(args)


text:
# -*- coding: utf-8 -*-
# @Date  : 2020/5/31
# @Author: Luokun
# @Email : olooook@outlook.com

import numpy as np

class LDA:
    """
    Linear Discriminant Analysis(线性判别分析)
    """
    ...

max_stars_repo_name: YangXin-Bio/ml

max_issues_repo_name: YangXin-Bio/ml

max_forks_repo_name: YangXin-Bio/ml


text:
import torch
from torch.autograd import Variable
from net import StackedHG
import numpy as np
from torchvision.datasets.folder import default_loader

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=-1)
parser.add_argument('--model')
parser.add_argument('--image')
args = parser.parse_args()

# model = StackedHG()
model = torch.load(args.model)
model.eval()

image = default_loader(args.image).astype(np.float32)
image = Variable(image, volatile=True)

# (1, 3, 256, 256) -> (1, 16, 64, 64)
heatmap = model([image])[0]
idx = np.argmax(heatmap.reshape(heatmap.shape[0], -1))
xx, yy = np.unravel_index(indices, (heatmap.shape[1:]))

keypoint = np.arange(heatmap.shape[0]), xx, yy
confidence = heatmap[np.arange(heatmap.shape[0]), xx, yy]

# visualize

# calculate MSE

max_stars_repo_name: fujibo/poseHG

max_issues_repo_name: fujibo/poseHG

max_forks_repo_name: fujibo/poseHG


text:
import os
from collections import OrderedDict
from numbers import Number

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import skimage

from softlearning.models.utils import flatten_input_structure
from .sac import SAC
from softlearning.samplers import rollouts
import math
from softlearning.misc.utils import save_video
from flatten_dict import flatten

tfd = tfp.distributions

def td_target(reward, discount, next_value):
    return reward + discount * next_value

class MultiSAC(SAC):
    """Soft Actor-Critic (SAC)

    References
    ----------
    [1] Tuomas Haarnoja*, Aurick Zhou*, Kristian Hartikainen*, George Tucker,
        Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter
        Abbeel, and Sergey Levine. Soft Actor-Critic Algorithms and
        Applications. arXiv preprint arXiv:1812.05905. 2018.
    """

    def __init__(
            self,
            training_environment,
            evaluation_environment,
            policies,
            Qs_per_policy,
            Q_targets_per_policy,
            pools,
            samplers,
            # goals=(),
            num_goals,
            plotter=None,

            # hyperparams shared across policies
            lr=3e-4,
            reward_scale=1.0,
            target_entropy='auto',
            discount=0.99,
            tau=5e-3,
            target_update_interval=1,
            action_prior='uniform',
            reparameterize=False,
            her_iters=0,
            save_full_state=False,
            save_eval_paths=False,
            per_alpha=1,
            normalize_ext_reward_gamma=1,
            ext_reward_coeffs=[],
            rnd_networks=(),
            rnd_lr=1e-4,
            rnd_int_rew_coeffs=[],
            rnd_gamma=0.99,
            save_reconstruction_frequency=1,
            **kwargs,
    ):
        """
        Args:
            env (`SoftlearningEnv`): Environment used for training.
            policy: A policy function approximator.
            initial_exploration_policy: ('Policy'): A policy that we use
                for initial exploration which is not trained by the algorithm.
            Qs: Q-function approximators. The min of these
                approximators will be used. Usage of at least two Q-functions
                improves performance by reducing overestimation bias.
            pool (`PoolBase`): Replay pool to add gathered samples to.
            plotter (`QFPolicyPlotter`): Plotter instance to be used for
                visualizing Q-function during training.
            lr (`float`): Learning rate used for the function approximators.
            discount (`float`): Discount factor for Q-function updates.
            tau (`float`): Soft value function target update weight.
            target_update_interval ('int'): Frequency at which target network
                updates occur in iterations.
            reparameterize ('bool'): If True, we use a gradient estimator for
                the policy derived using the reparameterization trick. We use
                a likelihood ratio based estimator otherwise.
        """

        super(SAC, self).__init__(**kwargs)

        self._training_environment = training_environment
        self._evaluation_environment = evaluation_environment
        assert hasattr(self._training_environment, "set_goal"), (
            "Need to implement `set_goal(goal_index)` method in training environment.")

        self._policies = policies
        self._Qs_per_policy = Qs_per_policy
        self._samplers = samplers
        self._pools = pools

        self._num_goals = num_goals
        self._goal_index = 0
        self._epoch_length *= num_goals
        self._n_epochs *= num_goals

        error_msg = 'Mismatch between number of policies, Qs, and samplers'
        assert len(self._policies) == num_goals, error_msg
        assert len(self._Qs_per_policy) == num_goals, error_msg
        assert len(self._samplers) == num_goals, error_msg

        self._Q_targets_per_policy = Q_targets_per_policy
        self._training_ops_per_policy = [{} for _ in range(num_goals)]

        self._policy_lr = lr
        self._Q_lr = lr

        self._reward_scale = reward_scale
        self._target_entropy = (
            -np.prod(self._training_environment.action_space.shape)
            if target_entropy == 'auto'
            else target_entropy)

        self._discount = discount
        self._tau = tau
        self._target_update_interval = target_update_interval
        self._action_prior = action_prior

        self._reparameterize = reparameterize
        self._plotter = plotter

        self._her_iters = her_iters
        self._base_env = training_environment.unwrapped

        self._save_full_state = save_full_state
        self._save_eval_paths = save_eval_paths

        self._n_episodes_elapsed = 0
        self._num_grad_steps_taken_per_policy = [0 for _ in range(self._num_goals)]

        self._normalize_ext_reward_gamma = normalize_ext_reward_gamma
        if ext_reward_coeffs:
            assert len(ext_reward_coeffs) == num_goals, (
                "Mismatch between number of goals and provided extrinsic reward coeffs")
            self._ext_reward_coeffs = ext_reward_coeffs
        else:
            self._ext_reward_coeffs = [1 for _ in range(num_goals)]
        self._running_ext_rew_stds = [1 for _ in range(num_goals)]
        self._rnd_int_rew_coeffs = []

        self._rnd_targets, self._rnd_predictors = None, None
        assert len(rnd_int_rew_coeffs) == num_goals, (
            "Mismatch between number of goals and provided intrinsic reward coeffs")
        self._rnd_int_rew_coeffs = rnd_int_rew_coeffs
        self._rnd_lr = rnd_lr
        self._rnd_gamma = rnd_gamma
        self._running_int_rew_stds = [1 for _ in range(num_goals)]
        if rnd_networks:
            assert len(rnd_networks) == num_goals
            self._rnd_targets = [rnd_network_pair[0] for rnd_network_pair in rnd_networks]
            self._rnd_predictors = [rnd_network_pair[1] for rnd_network_pair in rnd_networks]
        else:
            # RND reward coefficients must be 0 if there are no RND networks passed in
            for i in range(len(self._rnd_int_rew_coeffs)):
                self._rnd_int_rew_coeffs[i] = 0

        self._n_preprocessor_evals_per_epoch = 3
        self._online_vae = True
        self._save_reconstruction_frequency = save_reconstruction_frequency
        self._fixed_eval_pixels = None

        self._build()

    def _build(self):
        super(SAC, self)._build()
        self._init_external_rewards()
        if self._rnd_targets is not None and self._rnd_predictors is not None:
            self._init_rnd_updates()
        self._init_actor_updates()
        self._init_critic_updates()
        if self._uses_vae:
            self._init_vae_updates()
        elif self._uses_rae:
            self._init_rae_updates()
        self._init_diagnostics_ops()

    def _init_external_rewards(self):
        self._unscaled_ext_rewards = [
            self._placeholders['rewards'] for _ in range(self._num_goals)]

    def _get_Q_targets(self):
        Q_targets = []

        self._placeholders['reward'].update({
            f'running_ext_rew_std_{i}': tf.compat.v1.placeholder(
                tf.float32, shape=(), name=f'running_ext_rew_std_{i}')
            for i in range(self._num_goals)
        })

        (self._unscaled_int_rewards,
         self._int_rewards,
         self._normalized_ext_rewards,
         self._ext_rewards,
         self._total_rewards) = [], [], [], [], []

        for i, policy in enumerate(self._policies):
            policy_inputs = flatten_input_structure({
                name: self._placeholders['next_observations'][name]
                for name in policy.observation_keys
            })
            next_actions = policy.actions(policy_inputs)
            next_log_pis = policy.log_pis(policy_inputs, next_actions)

            next_Q_observations = {
                name: self._placeholders['next_observations'][name]
                for name in self._Qs_per_policy[i][0].observation_keys
            }
            next_Q_inputs = flatten_input_structure(
                {**next_Q_observations, 'actions': next_actions})
            next_Qs_values = tuple(Q(next_Q_inputs) for Q in self._Q_targets_per_policy[i])

            min_next_Q = tf.reduce_min(next_Qs_values, axis=0)
            next_values = min_next_Q - self._alphas[i] * next_log_pis

            terminals = tf.cast(self._placeholders['terminals'], next_values.dtype)

            if self._rnd_int_rew_coeffs[i] > 0:
                self._unscaled_int_rewards.append(tf.clip_by_value(
                    self._rnd_errors[i] / self._placeholders['reward'][f'running_int_rew_std_{i}'],
                    0, 1000
                ))
            else:
                self._unscaled_int_rewards.append(0)
            self._int_rewards.append(self._rnd_int_rew_coeffs[i] * self._unscaled_int_rewards[i])

            if self._ext_reward_coeffs[i] > 0:
                self._normalized_ext_rewards.append(
                    self._unscaled_ext_rewards[i] / self._placeholders['reward'][f'running_ext_rew_std_{i}'])
            else:
                self._normalized_ext_rewards.append(0)
            self._ext_rewards.append(self._ext_reward_coeffs[i] * self._normalized_ext_rewards[i])

            self._total_rewards.append(self._ext_rewards[i] + self._int_rewards[i])

            Q_target = td_target(
                reward=self._reward_scale * self._total_rewards[i],
                discount=self._discount,
                next_value=(1 - terminals) * next_values)
            Q_targets.append(tf.stop_gradient(Q_target))

        return Q_targets

    @property
    def _uses_vae(self):
        if self._policies and 'pixels' in self._policies[0].preprocessors:
            preprocessor = self._policies[0].preprocessors['pixels']
            return preprocessor.__class__.__name__ == 'OnlineVAEPreprocessor'
        return False

    @property
    def _uses_rae(self):
        if self._policies and 'pixels' in self._policies[0].preprocessors:
            preprocessor = self._policies[0].preprocessors['pixels']
            return preprocessor.__class__.__name__ == 'RAEPreprocessor'
        return False

    def _init_vae_updates(self):
        """
        Initializes VAE optimization update
        Creates a `tf.optimizer.minimize` operation, appends to `self._training_ops`.
        """
        # Need to loop through all the VAEs for Qs (2 per policy), and the VAEs for
        # the policies
        if not self._uses_vae:
            return

        vae_log_dir = os.path.join(os.getcwd(), 'vae')
        if not os.path.exists(vae_log_dir):
            os.makedirs(vae_log_dir)

        self._vae_optimizers_per_policy = []
        # The metrics will be used in `get_diagnostics` below
        self._vae_metrics_per_policy = []
        self._preprocessors_per_policy = []

        for i in range(self._num_goals):
            vae_metrics = {
                'latent_mean': {},
                'latent_logvar': {},
                'elbo': {},
                'reconstruction_loss': {},
                'kl_divergence': {},
            }
            if self._ext_reward_coeffs[i] != 0:
                # Only add update operations attached to preprocessors that are used
                # w.r.t the extrinsic reward
                policy_preprocessor = self._policies[i].preprocessors['pixels']
                Qs = self._Qs_per_policy[i]
                Q_preprocessor = Qs[0].observations_preprocessors['pixels']
                assert Q_preprocessor is Qs[1].observations_preprocessors['pixels'], (
                    'Preprocessors on the critics must be the same object.'
                )
                # preprocessors = (
                #     (f'policy_{i}', policy_preprocessor),
                #     (f'Q_{i}', Q_preprocessor),
                # )
                preprocessors = (
                    {f'policy_{i}': policy_preprocessor, f'Q_{i}': Q_preprocessor}
                    if policy_preprocessor is not Q_preprocessor
                    else {f'shared_{i}': policy_preprocessor}
                )
                self._preprocessors_per_policy.append(preprocessors)

                for name, vae in preprocessors.items():
                    # vae = preprocessor.vae
                    images = self._placeholders['observations']['pixels']
                    z_mean, z_log_var, z, reconstructions = vae(
                        images, include_reconstructions=True)

                    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
                        logits=reconstructions,
                        labels=tf.image.convert_image_dtype(images, tf.float32)
                    )
                    # z_mean, z_log_var, z = vae.encoder(images)

                    reconstruction_loss = tf.reduce_sum(cross_ent, axis=[1, 2, 3])
                    kl_divergence = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
                    kl_divergence = -0.5 * tf.reduce_sum(kl_divergence, axis=-1)

                    # ELBO = E(p(x|z)) - D_kl(q(z|x) || p(z)) * beta
                    elbo = -reconstruction_loss - vae.beta * kl_divergence
                    vae_metrics['latent_mean'][name] = tf.reduce_mean(z_mean)
                    vae_metrics['latent_logvar'][name] = tf.reduce_mean(z_log_var)
                    vae_metrics['elbo'][name] = tf.reduce_mean(elbo)
                    vae_metrics['reconstruction_loss'][name] = tf.reduce_mean(
                        reconstruction_loss)
                    vae_metrics['kl_divergence'][name] = tf.reduce_mean(kl_divergence)

                    vae_loss = -tf.reduce_mean(elbo)

                    vae_optimizer = tf.train.AdamOptimizer(
                        learning_rate=self._policy_lr,
                        name=f'{name}_vae_optimizer')

                    vae_train_op = vae_optimizer.minimize(
                        loss=vae_loss,
                        var_list=vae.trainable_variables
                    )
                    self._training_ops_per_policy[i].update({
                        f'{name}_vae_train_op': vae_train_op})
            # Add the metrics computed above
            self._vae_metrics_per_policy.append(vae_metrics)

    # TODO: Clean up these two updates, lots of redundancy
    def _init_rae_updates(self):
        """
        Initializes RAE optimization update
        Creates a `tf.optimizer.minimize` operation, appends to `self._training_ops`.
        """
        if not self._uses_rae:
            return

        rae_log_dir = os.path.join(os.getcwd(), 'rae')
        if not os.path.exists(rae_log_dir):
            os.makedirs(rae_log_dir)

        self._rae_optimizers_per_policy = []
        # The metrics will be used in `get_diagnostics` below
        self._rae_metrics_per_policy = []
        self._preprocessors_per_policy = []

        for i in range(self._num_goals):
            rae_metrics = {
                'reconstruction_loss': {},
                'latent_l2_norm': {},
            }
            if self._ext_reward_coeffs[i] != 0:
                # Only add update operations attached to preprocessors that are used
                # w.r.t the extrinsic reward
                policy_preprocessor = self._policies[i].preprocessors['pixels']
                Qs = self._Qs_per_policy[i]
                Q_preprocessor = Qs[0].observations_preprocessors['pixels']
                assert Q_preprocessor is Qs[1].observations_preprocessors['pixels'], (
                    'Preprocessors on the critics must be the same object.'
                )
                # preprocessors = (
                #     (f'policy_{i}', policy_preprocessor),
                #     (f'Q_{i}', Q_preprocessor),
                # )
                preprocessors = (
                    {f'policy_{i}': policy_preprocessor, f'Q_{i}': Q_preprocessor}
                    if policy_preprocessor is not Q_preprocessor
                    else {f'shared_{i}': policy_preprocessor}
                )
                self._preprocessors_per_policy.append(preprocessors)

                for name, rae in preprocessors.items():
                    # rae = preprocessor.rae

                    images = self._placeholders['observations']['pixels']
                    z = rae.encoder(images)
                    reconstructions = rae.decoder(z)

                    # L_RAE = ||true - reconstructions||^2 + ||z||^2 + ||theta||^2
                    # Sum over individual pixel MSE over the entire image
                    reconstruction_loss = tf.reduce_mean(tf.reduce_sum(
                        tf.square(
                            tf.image.convert_image_dtype(images, tf.float32) - reconstructions
                        ), axis=[1, 2, 3]))
                    latent_l2_norm_loss = tf.reduce_mean(tf.reduce_sum(tf.square(z), axis=-1))
                    # latent_l2_norm_loss = tf.reduce_mean(tf.norm(z, axis=-1))
                    rae_loss = reconstruction_loss + latent_l2_norm_loss

                    rae_metrics['reconstruction_loss'][name] = reconstruction_loss
                    rae_metrics['latent_l2_norm'][name] = latent_l2_norm_loss

                    rae_optimizer = tf.train.AdamOptimizer(
                        learning_rate=self._policy_lr,
                        name=f'{name}_rae_optimizer')

                    rae_train_op = rae_optimizer.minimize(
                        loss=rae_loss,
                        var_list=rae.trainable_variables
                    )
                    self._training_ops_per_policy[i].update({
                        f'{name}_rae_train_op': rae_train_op})
            # Add the metrics computed above
            self._rae_metrics_per_policy.append(rae_metrics)

    def _init_critic_updates(self):
        """Create minimization operation for critics' Q-functions.

        Creates a `tf.optimizer.minimize` operation for updating
        critic Q-function with gradient descent, and appends it to
        `self._training_ops` attribute.

        See Equations (5, 6) in [1], for further information of the
        Q-function update rule.
        """
        Q_targets = self._get_Q_targets()
        assert len(Q_targets) == len(self._policies)
        for Q_target in Q_targets:
            assert Q_target.shape.as_list() == [None, 1]

        self._Q_optimizers_per_policy = []
        self._Q_values_per_policy = []
        self._Q_losses_per_policy = []

        for i, Qs in enumerate(self._Qs_per_policy):
            Q_observations = {
                name: self._placeholders['observations'][name]
                for name in Qs[0].observation_keys
            }
            Q_inputs = flatten_input_structure({
                **Q_observations, 'actions': self._placeholders['actions']})

            Q_values = tuple(Q(Q_inputs) for Q in Qs)
            self._Q_values_per_policy.append(Q_values)

            Q_losses = tuple(
                tf.compat.v1.losses.mean_squared_error(
                    labels=Q_targets[i], predictions=Q_value, weights=0.5)
                for Q_value in Q_values)
            self._Q_losses_per_policy.append(Q_losses)

            # self._bellman_errors.append(tf.reduce_min(tuple(
            #     tf.math.squared_difference(Q_target, Q_value)
            #     for Q_value in Q_values), axis=0))

            Q_optimizers = tuple(
                tf.compat.v1.train.AdamOptimizer(
                    learning_rate=self._Q_lr,
                    name='{}_{}_optimizer_{}'.format(i, Q._name, j)
                ) for j, Q in enumerate(Qs))
            self._Q_optimizers_per_policy.append(Q_optimizers)

            Q_training_ops = tuple(
                Q_optimizer.minimize(loss=Q_loss, var_list=Q.trainable_variables)
                for i, (Q, Q_loss, Q_optimizer)
                in enumerate(zip(Qs, Q_losses, Q_optimizers)))

            self._training_ops_per_policy[i].update({f'Q_{i}': tf.group(Q_training_ops)})

    def _init_actor_updates(self):
        """Create minimization operations for policies and entropies.

        Creates a `tf.optimizer.minimize` operations for updating
        policy and entropy with gradient descent, and adds them to
        `self._training_ops` attribute.

        See Section 4.2 in [1], for further information of the policy update,
        and Section 5 in [1] for further information of the entropy update.
        """

        self._log_alphas = []
        self._alpha_optimizers = []
        self._alphas = []
        self._policy_optimizers = []
        self._policy_losses = []

        for i, policy in enumerate(self._policies):
            policy_inputs = flatten_input_structure({
                name: self._placeholders['observations'][name]
                for name in policy.observation_keys
            })
            actions = policy.actions(policy_inputs)
            log_pis = policy.log_pis(policy_inputs, actions)

            assert log_pis.shape.as_list() == [None, 1]

            log_alpha = tf.compat.v1.get_variable(
                f'log_alpha_{i}',
                dtype=tf.float32,
                initializer=0.0)
            alpha = tf.exp(log_alpha)
            self._log_alphas.append(log_alpha)
            self._alphas.append(alpha)

            if isinstance(self._target_entropy, Number):
                alpha_loss = -tf.reduce_mean(
                    log_alpha * tf.stop_gradient(log_pis + self._target_entropy))

                alpha_optimizer = tf.compat.v1.train.AdamOptimizer(
                    self._policy_lr, name=f'alpha_optimizer_{i}')
                self._alpha_optimizers.append(alpha_optimizer)
                alpha_train_op = alpha_optimizer.minimize(loss=alpha_loss, var_list=[log_alpha])
                self._training_ops_per_policy[i].update({
                    f'temperature_alpha_{i}': alpha_train_op
                })

            if self._action_prior == 'normal':
                policy_prior = tfp.distributions.MultivariateNormalDiag(
                    loc=tf.zeros(self._action_shape),
                    scale_diag=tf.ones(self._action_shape))
                policy_prior_log_probs = policy_prior.log_prob(actions)
            elif self._action_prior == 'uniform':
                policy_prior_log_probs = 0.0

            Q_observations = {
                name: self._placeholders['observations'][name]
                for name in self._Qs_per_policy[i][0].observation_keys
            }
            Q_inputs = flatten_input_structure({
                **Q_observations, 'actions': actions})
            Q_log_targets = tuple(Q(Q_inputs) for Q in self._Qs_per_policy[i])
            min_Q_log_target = tf.reduce_min(Q_log_targets, axis=0)

            if self._reparameterize:
                policy_kl_losses = (
                    alpha * log_pis
                    - min_Q_log_target
                    - policy_prior_log_probs)
            else:
                raise NotImplementedError

            assert policy_kl_losses.shape.as_list() == [None, 1]

            self._policy_losses.append(policy_kl_losses)
            policy_loss = tf.reduce_mean(policy_kl_losses)

            policy_optimizer = tf.compat.v1.train.AdamOptimizer(
                learning_rate=self._policy_lr,
                name=f"policy_optimizer_{i}")

            self._policy_optimizers.append(policy_optimizer)

            policy_train_op = policy_optimizer.minimize(
                loss=policy_loss,
                var_list=policy.trainable_variables)

            self._training_ops_per_policy[i].update({f'policy_train_op_{i}': policy_train_op})

    def _init_rnd_updates(self):
        (self._rnd_errors,
         self._rnd_losses,
         self._rnd_error_stds,
         self._rnd_optimizers) = [], [], [], []
        for i in range(self._num_goals):
            self._placeholders['reward'].update({
                f'running_int_rew_std_{i}': tf.compat.v1.placeholder(
                    tf.float32, shape=(), name=f'running_int_rew_std_{i}')
            })
            policy_inputs = flatten_input_structure({
                name: self._placeholders['observations'][name]
                for name in self._rnd_predictors[i].observation_keys
            })

            targets = tf.stop_gradient(self._rnd_targets[i](policy_inputs))
            predictions = self._rnd_predictors[i](policy_inputs)

            self._rnd_errors.append(tf.expand_dims(tf.reduce_mean(
                tf.math.squared_difference(targets, predictions), axis=-1), 1))
            self._rnd_losses.append(tf.reduce_mean(self._rnd_errors[i]))
            self._rnd_error_stds.append(tf.math.reduce_std(self._rnd_errors[i]))
            self._rnd_optimizers.append(tf.compat.v1.train.AdamOptimizer(
                learning_rate=self._rnd_lr,
                name=f"rnd_optimizer_{i}"))
            rnd_train_op = self._rnd_optimizers[i].minimize(
                loss=self._rnd_losses[i])
            self._training_ops_per_policy[i].update(
                {f'rnd_train_op_{i}': rnd_train_op}
            )

    def _init_diagnostics_ops(self):
        diagnosables_per_goal = [
            OrderedDict((
                (f'Q_value_{i}', self._Q_values_per_policy[i]),
                (f'Q_loss_{i}', self._Q_losses_per_policy[i]),
                (f'policy_loss_{i}', self._policy_losses[i]),
                (f'alpha_{i}', self._alphas[i])
            ))
            for i in range(self._num_goals)
        ]

        for i in range(self._num_goals):
            # Only record the intrinsic/extrinsic reward diagnostics if
            # the reward is actually used (i.e. the reward coeff is not 0)
            if self._rnd_int_rew_coeffs[i]:
                diagnosables_per_goal[i][f'rnd_reward_{i}'] = self._int_rewards[i]
                diagnosables_per_goal[i][f'rnd_error_{i}'] = self._rnd_errors[i]
                diagnosables_per_goal[i][f'running_rnd_reward_std_{i}'] = (
                    self._placeholders['reward'][f'running_int_rew_std_{i}'])

            if self._ext_reward_coeffs[i]:
                diagnosables_per_goal[i][f'ext_reward_{i}'] = self._ext_rewards[i]
                diagnosables_per_goal[i][f'normalized_ext_reward_{i}'] = (
                    self._normalized_ext_rewards[i])
                diagnosables_per_goal[i][f'unnormalized_ext_reward_{i}'] = (
                    self._unscaled_ext_rewards[i])

            diagnosables_per_goal[i][f'running_ext_reward_std_{i}'] = (
                self._placeholders['reward'][f'running_ext_rew_std_{i}'])
            diagnosables_per_goal[i][f'total_reward_{i}'] = self._total_rewards[i]

        if self._uses_vae:
            for i in range(self._num_goals):
                vae_metrics = self._vae_metrics_per_policy[i]
                for metric, logs in vae_metrics.items():
                    for name, log in logs.items():
                        diagnostic_key = f'vae/{name}/{metric}'
                        diagnosables_per_goal[i][diagnostic_key] = log
        elif self._uses_rae:
            for i in range(self._num_goals):
                rae_metrics = self._rae_metrics_per_policy[i]
                for metric, logs in rae_metrics.items():
                    for name, log in logs.items():
                        diagnostic_key = f'rae/{name}/{metric}'
                        diagnosables_per_goal[i][diagnostic_key] = log

        diagnostic_metrics = OrderedDict((
            ('mean', tf.reduce_mean),
            ('std', lambda x: tfp.stats.stddev(x, sample_axis=None)),
            ('max', tf.math.reduce_max),
            ('min', tf.math.reduce_min),
        ))

        self._diagnostics_ops_per_goal = [
            OrderedDict([
                (f'{key}-{metric_name}', metric_fn(values))
                for key, values in diagnosables.items()
                for metric_name, metric_fn in diagnostic_metrics.items()
            ])
            for diagnosables in diagnosables_per_goal
        ]

    def _training_batch(self, batch_size=None):
        return self._samplers[self._goal_index].random_batch(batch_size)

    def _evaluation_batches(self, batch_size=None):
        return [self._samplers[i].random_batch(batch_size) for i in range(self._num_goals)]

    def _update_target(self, i, tau=None):
        """ Update target networks for policy i. """
        tau = tau or self._tau

        for Q, Q_target in zip(self._Qs_per_policy[i], self._Q_targets_per_policy[i]):
            source_params = Q.get_weights()
            target_params = Q_target.get_weights()
            Q_target.set_weights([
                tau * source + (1.0 - tau) * target
                for source, target in zip(source_params, target_params)
            ])

    def _epoch_before_hook(self, *args, **kwargs):
        super(SAC, self)._epoch_before_hook(*args, **kwargs)

    def _do_training(self, iteration, batch):
        """Runs the operations for updating training and target ops."""
        feed_dict = self._get_feed_dict(iteration, batch)

        self._session.run(self._training_ops_per_policy[self._goal_index], feed_dict)

        if self._rnd_int_rew_coeffs[self._goal_index]:
            int_rew_std = np.maximum(np.std(self._session.run(
                self._unscaled_int_rewards[self._goal_index], feed_dict)), 1e-3)
            self._running_int_rew_stds[
                self._goal_index] = self._running_int_rew_stds[self._goal_index] * self._rnd_gamma + int_rew_std * (1-self._rnd_gamma)

        if self._normalize_ext_reward_gamma != 1 and self._ext_reward_coeffs[self._goal_index]:
            ext_rew_std = np.maximum(np.std(self._session.run(
                self._normalized_ext_rewards[self._goal_index], feed_dict)), 1e-3)
            self._running_ext_rew_stds[
                self._goal_index] = self._running_ext_rew_stds[self._goal_index] * self._normalize_ext_reward_gamma + \
                ext_rew_std * (1-self._normalize_ext_reward_gamma)

        self._num_grad_steps_taken_per_policy[self._goal_index] += 1

        if self._num_grad_steps_taken_per_policy[self._goal_index] % self._target_update_interval == 0:
            # Run target ops here.
            self._update_target(self._goal_index)

    def _get_feed_dict(self, iteration, batch):
        batch_flat = flatten(batch)
        placeholders_flat = flatten(self._placeholders)

        # if np.random.rand() < 1e-4 and 'pixels' in batch['observations']:
        #     import os
        #     from skimage import io
        #     random_idx = np.random.randint(
        #         batch['observations']['pixels'].shape[0])
        #     image_save_dir = os.path.join(os.getcwd(), 'pixels')
        #     image_save_path = os.path.join(
        #         image_save_dir, f'observation_{iteration}_batch.png')
        #     if not os.path.exists(image_save_dir):
        #         os.makedirs(image_save_dir)
        #     io.imsave(image_save_path,
        #               batch['observations']['pixels'][random_idx].copy())

        feed_dict = {
            placeholders_flat[key]: batch_flat[key]
            for key in placeholders_flat.keys()
            if key in batch_flat.keys()
        }

        if iteration is not None:
            feed_dict[self._placeholders['iteration']] = iteration

        ext_rew_std_ph = self._placeholders['reward'][f'running_ext_rew_std_{self._goal_index}']
        feed_dict[ext_rew_std_ph] = self._running_ext_rew_stds[self._goal_index]
        if self._rnd_int_rew_coeffs[self._goal_index]:
            int_rew_std_ph = self._placeholders['reward'][f'running_int_rew_std_{self._goal_index}']
            feed_dict[int_rew_std_ph] = self._running_int_rew_stds[self._goal_index]

        return feed_dict

    def get_diagnostics(self,
                        iteration,
                        batches,
                        training_paths_per_policy,
                        evaluation_paths_per_policy):
        """Return diagnostic information as ordered dictionary.

        Also calls the `draw` method of the plotter, if plotter defined.
        """
        goal_index = self._goal_index
        diagnostics = {}

        for i in range(self._num_goals):
            self._goal_index = i
            feed_dict = self._get_feed_dict(iteration, batches[i])
            diagnostics.update(
                self._session.run({**self._diagnostics_ops_per_goal[i]}, feed_dict))
            diagnostics.update(OrderedDict([
                (f'policy_{i}/{key}', value)
                for key, value in self._policies[i].get_diagnostics(
                    flatten_input_structure({
                        name: batches[i]['observations'][name]
                        for name in self._policies[i].observation_keys})
                ).items()
            ]))
        self._goal_index = goal_index
        should_save = iteration % self._save_reconstruction_frequency == 0

        # Generate random pixels to evaluate the preprocessors
        if 'pixels' in self._placeholders['observations']:
            random_idxs = np.random.choice(
                feed_dict[self._placeholders['observations']['pixels']].shape[0],
                size=self._n_preprocessor_evals_per_epoch)
            eval_pixels = (
                feed_dict[self._placeholders['observations']['pixels']][random_idxs])
        else:
            eval_pixels = None

        if self._uses_vae and should_save:
            assert eval_pixels
            for i, preprocessors in enumerate(self._preprocessors_per_policy):
                if self._ext_reward_coeffs[i] == 0:
                    continue

                for name, vae in preprocessors.items():
                    z_mean, z_logvar, z = self._session.run(vae.encoder(eval_pixels))
                    reconstructions = self._session.run(
                        tf.math.sigmoid(vae.decoder(z)))
                    concat = np.concatenate([
                        eval_pixels,
                        skimage.util.img_as_ubyte(reconstructions)
                    ], axis=2)

                    sampled_z = np.random.normal(
                        size=(eval_pixels.shape[0], vae.latent_dim))
                    decoded_samples = self._session.run(
                        tf.math.sigmoid(vae.decoder.output),
                        feed_dict={vae.decoder.input: sampled_z}
                    )

                    save_path = os.path.join(os.getcwd(), 'vae')
                    recon_concat = np.vstack(concat)
                    skimage.io.imsave(
                        os.path.join(
                            save_path,
                            f'{name}_reconstruction_{iteration}.png'),
                        recon_concat)
                    samples_concat = np.vstack(decoded_samples)
                    skimage.io.imsave(
                        os.path.join(
                            save_path,
                            f'{name}_sample_{iteration}.png'),
                        samples_concat)
        elif self._uses_rae:
            assert eval_pixels
            for i, preprocessors in enumerate(self._preprocessors_per_policy):
                if self._ext_reward_coeffs[i] == 0:
                    continue
                for name, rae in preprocessors.items():
                    z, reconstructions = self._session.run(
                        rae(eval_pixels, include_reconstructions=True))
                    concat = np.concatenate([
                        eval_pixels,
                        skimage.util.img_as_ubyte(reconstructions)
                    ], axis=2)

                    if should_save:
                        save_path = os.path.join(os.getcwd(), 'rae')
                        recon_concat = np.vstack(concat)
                        skimage.io.imsave(
                            os.path.join(
                                save_path,
                                f'{name}_reconstruction_{iteration}.png'),
                            recon_concat)

                    # Track latents
                    if self._fixed_eval_pixels is None:
                        self._fixed_eval_pixels = eval_pixels
                        self._fixed_eval_latents = np.zeros(z.shape)

                    z_fixed, reconstructions_fixed = self._session.run(
                        rae(self._fixed_eval_pixels, include_reconstructions=True))

                    if should_save:
                        concat_fixed = np.concatenate([
                            self._fixed_eval_pixels,
                            skimage.util.img_as_ubyte(reconstructions_fixed)
                        ], axis=2)
                        recon_concat_fixed = np.vstack(concat_fixed)
                        skimage.io.imsave(
                            os.path.join(
                                save_path,
                                f'{name}_fixed_reconstruction_{iteration}.png'),
                            recon_concat_fixed)

                    z_diff = np.linalg.norm(z_fixed - self._fixed_eval_latents, axis=1)
                    diagnostics.update({
                        f'rae/{name}/tracked-latent-l2-difference-with-prev-mean': np.mean(z_diff),
                        f'rae/{name}/tracked-latent-l2-difference-with-prev-std': np.std(z_diff),
                    })
                    # Save the previous latents to compare in the next epoch
                    self._fixed_eval_latents = z_fixed

        if self._save_eval_paths:
            import pickle
            file_name = f'eval_paths_{iteration // self.epoch_length}.pkl'
            with open(os.path.join(os.getcwd(), file_name)) as f:
                pickle.dump(evaluation_paths_per_policy, f)

        if self._plotter:
            self._plotter.draw()

        return diagnostics

    @property
    def tf_saveables(self):
        saveables = {
            **{
                f'_policy_optimizer_{i}': policy_optimizer
                for i, policy_optimizer in enumerate(self._policy_optimizers)
            },
            **{
                f'_log_alphas_{i}': log_alpha
                for i, log_alpha in enumerate(self._log_alphas)
            },
        }

        Q_optimizer_saveables = [
            {
                f'Q_optimizer_{i}_{j}': Q_optimizer
                for j, Q_optimizer in enumerate(Q_optimizers)
            }
            for i, Q_optimizers in enumerate(self._Q_optimizers_per_policy)
        ]

        for Q_opt_dict in Q_optimizer_saveables:
            saveables.update(Q_opt_dict)

        if hasattr(self, '_alpha_optimizer'):
            saveables['_alpha_optimizer'] = self._alpha_optimizer

        return saveables

    def _initial_exploration_hook(self, env, initial_exploration_policy, goal_index):
        print("start random exploration")
        if self._n_initial_exploration_steps < 1:
            return

        if not initial_exploration_policy:
            raise ValueError(
                "Initial exploration policy must be provided when"
                " n_initial_exploration_steps > 0.")
        self._set_goal(goal_index)
        # env.set_goal(goal_index)

        self._samplers[goal_index].initialize(env, initial_exploration_policy, self._pools[goal_index])
        while self._pools[goal_index].size < self._n_initial_exploration_steps:
            self._samplers[goal_index].sample()

    def _init_training(self):
        for i in range(self._num_goals):
            self._update_target(i, tau=1.0)

    def _initialize_samplers(self):
        for i, sampler in enumerate(self._samplers):
            sampler.initialize(self._training_environment, self._policies[i], self._pools[i])
            sampler.set_save_training_video_frequency(self._save_training_video_frequency)
            if hasattr(sampler, 'set_algorithm'):
                sampler.set_algorithm(self)
        self._n_episodes_elapsed = sum([self._samplers[i]._n_episodes for i in range(self._num_goals)])

    @property
    def ready_to_train(self):
        return self._samplers[self._goal_index].batch_ready()

    def _do_sampling(self, timestep):
        self._sample_count += 1
        self._samplers[self._goal_index].sample()

    def _set_goal(self, goal_index):
        """ Set goal in env. """
        assert goal_index >= 0 and goal_index < self._num_goals
        # print("setting goal to: ", goal_index, ", n_episodes_elapsed: ", self._n_episodes_elapsed)
        self._goal_index = goal_index
        self._training_environment.set_goal(self._goal_index) # TODO: implement in env

    def _training_before_hook(self, *args, **kwargs):
        super()._training_before_hook(*args, **kwargs)
        # Reset goal
        self._set_goal(0)

    def _timestep_before_hook(self, *args, **kwargs):
        # Check if an entire trajectory has been completed
        n_episodes_sampled = sum([self._samplers[i]._n_episodes for i in range(self._num_goals)])
        if n_episodes_sampled > self._n_episodes_elapsed:
            self._n_episodes_elapsed = n_episodes_sampled
            new_goal_index = (self._goal_index + 1) % self._num_goals
            self._set_goal(new_goal_index)

    def _train(self):
        """Return a generator that performs RL training.

        Args:
            env (`SoftlearningEnv`): Environment used for training.
            policy (`Policy`): Policy used for training
            initial_exploration_policy ('Policy'): Policy used for exploration
                If None, then all exploration is done using policy
            pool (`PoolBase`): Sample pool to add samples to
        """
        import gtimer as gt
        from itertools import count
        training_environment = self._training_environment
        evaluation_environment = self._evaluation_environment
        training_metrics = [0 for _ in range(self._num_goals)]

        if not self._training_started:
            self._init_training()

            for i in range(self._num_goals):
                self._initial_exploration_hook(
                    training_environment, self._initial_exploration_policy, i)

        self._initialize_samplers()
        self._sample_count = 0

        gt.reset_root()
        gt.rename_root('RLAlgorithm')
        gt.set_def_unique(False)

        print("starting_training")
        self._training_before_hook()
        import time

        for self._epoch in gt.timed_for(range(self._epoch, self._n_epochs)):
            self._epoch_before_hook()
            gt.stamp('epoch_before_hook')
            start_samples = sum([self._samplers[i]._total_samples for i in range(self._num_goals)])
            sample_times = []
            for i in count():
                samples_now = sum([self._samplers[i]._total_samples for i in range(self._num_goals)])
                self._timestep = samples_now - start_samples

                if samples_now >= start_samples + self._epoch_length and self.ready_to_train:
                    break

                t0 = time.time()
                self._timestep_before_hook()
                gt.stamp('timestep_before_hook')

                self._do_sampling(timestep=self._total_timestep)
                gt.stamp('sample')
                sample_times.append(time.time() - t0)
                t0 = time.time()
                if self.ready_to_train:
                    self._do_training_repeats(timestep=self._total_timestep)
                gt.stamp('train')
                # print("Train time: ", time.time() - t0)

                self._timestep_after_hook()
                gt.stamp('timestep_after_hook')

            # TODO diagnostics per goal
            print("Average Sample Time: ", np.mean(np.array(sample_times)))
            print("Step count", self._sample_count)
            training_paths_per_policy = self._training_paths()
            # self.sampler.get_last_n_paths(
            #     math.ceil(self._epoch_length / self.sampler._max_path_length))
            gt.stamp('training_paths')
            evaluation_paths_per_policy = self._evaluation_paths()
            gt.stamp('evaluation_paths')

            training_metrics_per_policy = self._evaluate_rollouts(
                training_paths_per_policy, training_environment)
            gt.stamp('training_metrics')

            if evaluation_paths_per_policy:
                evaluation_metrics_per_policy = self._evaluate_rollouts(
                    evaluation_paths_per_policy, evaluation_environment)
                gt.stamp('evaluation_metrics')
            else:
                evaluation_metrics_per_policy = [{} for _ in range(self._num_goals)]

            self._epoch_after_hook(training_paths_per_policy)
            gt.stamp('epoch_after_hook')

            t0 = time.time()

            sampler_diagnostics_per_policy = [
                self._samplers[i].get_diagnostics() for i in range(self._num_goals)]

            diagnostics = self.get_diagnostics(
                iteration=self._total_timestep,
                batches=self._evaluation_batches(),
                training_paths_per_policy=training_paths_per_policy,
                evaluation_paths_per_policy=evaluation_paths_per_policy)

            time_diagnostics = gt.get_times().stamps.itrs

            print("Basic diagnostics: ", time.time() - t0)
            print("Sample count: ", self._sample_count)

            diagnostics.update(OrderedDict((
                *(
                    (f'times/{key}', time_diagnostics[key][-1])
                    for key in sorted(time_diagnostics.keys())
                ),
                ('epoch', self._epoch),
                ('timestep', self._timestep),
                ('timesteps_total', self._total_timestep),
                ('train-steps', self._num_train_steps),
            )))

            print("Other basic diagnostics: ", time.time() - t0)
            for i, (evaluation_metrics, training_metrics, sampler_diagnostics) in (
                enumerate(zip(evaluation_metrics_per_policy,
                              training_metrics_per_policy,
                              sampler_diagnostics_per_policy))):
                diagnostics.update(OrderedDict((
                    *(
                        (f'evaluation_{i}/{key}', evaluation_metrics[key])
                        for key in sorted(evaluation_metrics.keys())
                    ),
                    *(
                        (f'training_{i}/{key}', training_metrics[key])
                        for key in sorted(training_metrics.keys())
                    ),
                    *(
                        (f'sampler_{i}/{key}', sampler_diagnostics[key])
                        for key in sorted(sampler_diagnostics.keys())
                    ),
                )))

            # if self._eval_render_kwargs and hasattr(
            #         evaluation_environment, 'render_rollouts'):
            #     # TODO(hartikainen): Make this consistent such that there's no
            #     # need for the hasattr check.
            #     training_environment.render_rollouts(evaluation_paths)

            yield diagnostics
            print("Diagnostic time: ",  time.time() - t0)

        for i in range(self._num_goals):
            self._samplers[i].terminate()

        self._training_after_hook()

        del evaluation_paths_per_policy

        yield {'done': True, **diagnostics}

    def _training_paths(self):
        """ Override to interleave training videos between policy rollouts. """
        num_paths = math.ceil(
            (self._epoch_length / self._num_goals) / self._samplers[self._goal_index]._max_path_length)
        paths_per_policy = [
            self._samplers[i].get_last_n_paths(num_paths)
            for i in range(self._num_goals)
        ]

        if (self._save_training_video_frequency > 0
            and self._epoch % self._save_training_video_frequency == 0):
            i = 0
            for rollout_num in reversed(range(len(paths_per_policy[0]))):
                video_frames = []
                for goal_index in range(self._num_goals):
                    video_frames.extend(paths_per_policy[goal_index][rollout_num].pop('images'))
                video_file_name = f'training_path_{self._epoch}_{i}.mp4'
                video_file_path = os.path.join(
                    os.getcwd(), 'videos', video_file_name)
                fps = 60
                # fps = 1 // getattr(self._training_environment, 'dt', 1/30)
                save_video(video_frames, video_file_path, fps=fps)
                i += 1

        return paths_per_policy

    def _evaluation_paths(self):
        if self._eval_n_episodes < 1:
            return ()

        should_save_video = (
            self._video_save_frequency > 0
            and self._epoch % self._video_save_frequency == 0)

        paths = []
        for goal_index in range(self._num_goals):
            with self._policies[goal_index].set_deterministic(self._eval_deterministic):
                self._evaluation_environment.set_goal(goal_index)
                paths.append(
                    rollouts(
                        self._eval_n_episodes,
                        self._evaluation_environment,
                        self._policies[goal_index],
                        self._samplers[goal_index]._max_path_length,
                        render_kwargs=(self._eval_render_kwargs
                                       if should_save_video else {})
                    )
                )

        # TODO: interleave videos from different policies
        if should_save_video:
            fps = 1 // getattr(self._evaluation_environment, 'dt', 1/30)
            for rollout_num in range(len(paths[0])):
                video_frames = []
                for goal_index in range(self._num_goals):
                    video_frames.append(paths[goal_index][rollout_num].pop('images'))
                video_frames = np.concatenate(video_frames)
                video_file_name = f'evaluation_path_{self._epoch}_{rollout_num}.mp4'
                video_file_path = os.path.join(
                    os.getcwd(), 'videos', video_file_name)
                save_video(video_frames, video_file_path, fps=fps)
        return paths

    def _evaluate_rollouts(self, episodes_per_policy, env):
        """Compute evaluation metrics for the given rollouts."""
        diagnostics_per_policy = []
        for i, episodes in enumerate(episodes_per_policy):
            episodes_rewards = [episode['rewards'] for episode in episodes]
            episodes_reward = [np.sum(episode_rewards)
                               for episode_rewards in episodes_rewards]
            episodes_length = [episode_rewards.shape[0]
                               for episode_rewards in episodes_rewards]

            diagnostics = OrderedDict((
                ('episode-reward-mean', np.mean(episodes_reward)),
                ('episode-reward-min', np.min(episodes_reward)),
                ('episode-reward-max', np.max(episodes_reward)),
                ('episode-reward-std', np.std(episodes_reward)),
                ('episode-length-mean', np.mean(episodes_length)),
                ('episode-length-min', np.min(episodes_length)),
                ('episode-length-max', np.max(episodes_length)),
                ('episode-length-std', np.std(episodes_length)),
            ))

            env_infos = env.get_path_infos(episodes)
            for key, value in env_infos.items():
                diagnostics[f'env_infos/{key}'] = value

            diagnostics_per_policy.append(diagnostics)

        return diagnostics_per_policy

max_stars_repo_name: facebookresearch/MTRF

max_issues_repo_name: facebookresearch/MTRF

max_forks_repo_name: facebookresearch/MTRF


text:
"""Routines to help draw star charts."""

import numpy as np
from IPython.display import HTML
from matplotlib import rc_context, rcParams
from matplotlib.animation import FuncAnimation, FFMpegWriter, writers
from .starlib import Star

def _plot_stars(catalog, observer, project, ax, mag1, mag2, margin=1.25):
    """Experiment in progress, hence the underscore; expect changes."""

    art = []

    # from astropy import wcs
    # w = wcs.WCS(naxis=2)
    # w.wcs.crpix = [-234.75, 8.3393]
    # w.wcs.cdelt = np.array([-0.066667, 0.066667])
    # w.wcs.crval = [0, -90]
    # w.wcs.ctype = ["RA---AIR", "DEC--AIR"]
    # w.wcs.set_pv([(2, 1, 45.0)])

    # import matplotlib.pyplot as plt

    # plt.subplot(projection=wcs)
    # #plt.imshow(hdu.data, vmin=-2.e-5, vmax=2.e-4, origin='lower')
    # plt.grid(color='white', ls='solid')
    # plt.xlabel('Galactic Longitude')
    # plt.ylabel('Galactic Latitude')

    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_xlim()
    lim = max(abs(xmin), abs(xmax), abs(ymin), abs(ymax)) * margin
    lims = (-lim, lim)
    ax.set_xlim(lims)
    ax.set_ylim(lims)
    ax.set_aspect('equal')

    o = observer[0]

    # Dim stars: points of with varying gray levels.

    c = catalog
    c = c[c['magnitude'] > mag1]
    c = c[c['magnitude'] <= mag2]
    #print('Second star group:', len(c))
    c = c.sort_values('magnitude', ascending=False)
    s = Star(ra_hours=c.ra_hours, dec_degrees=c.dec_degrees)
    spos = o.observe(s)
    x, y = project(spos)
    m = (mag2 - c['magnitude']) / (mag2 - mag1)
    # Note that "gray_r" is white for 0.0 and black for 1.0
    art.append(ax.scatter(
        x, y, s=1.0,
        c=1 - 0.8 * m, cmap='gray_r', vmin=0.0, vmax=1.0,
    ))

    # Bright stars: black circles of varying radius, surrounded by a
    # white gap in case stars are touching.  Draw the brightest stars
    # first to stop them from completely occluding smaller companions.

    def mag_to_radius(m):
        return (mag1 - m) * scale + 1.0

    c = catalog
    c = c[c['magnitude'] <= mag1]
    c = c.sort_values('magnitude', ascending=True)
    #print('First star group:', len(c))
    s = Star(ra_hours=c.ra_hours, dec_degrees=c.dec_degrees)
    spos = o.observe(s)
    x, y = project(spos)
    scale = 1.5
    radius = mag_to_radius(c['magnitude'])

    x2 = np.repeat(x, 2)
    y2 = np.repeat(y, 2)
    radius2 = (radius[:,None] + (3.0,0.0)).flatten()
    c2 = ('w', 'k')
    c2 = ('k', 'w')

    art.append(ax.scatter(x2, y2, s=radius2 ** 2.0, c=c2))

    return art, mag_to_radius

@writers.register('ffmpeg_small')
class _FFMpegWriter(FFMpegWriter):
    def _args(self):
        args = super(_FFMpegWriter, self)._args()
        i = args.index('h264') + 1
        args[i:i] = ['-x264-params', 'keyint=99999999:scenecut=0']
        #args.append('keyint=123:min-keyint=20')
        return args

X = []

class _Animation(FuncAnimation):
    patcher = None

    def _draw_next_frame(self, framedata, blit):
        blit = True    # override matplotlib's refusal to blit in save()
        super(_Animation, self)._draw_next_frame(framedata, blit)

        if self.patcher is None:
            self.patcher = patch('matplotlib.backends.backend_agg'
                            '.FigureCanvasAgg.draw', draw)
            self.patcher.__enter__()

    # Experiment to avoid drawing month names over each other when they
    # are outside the axis borders and don't get blitted:

    # def _pre_draw(self, framedata, blit):
    #     artists = self._drawn_artists
    #     for figure in set(a.axes.figure for a in artists):
    #         from IPython.core.debugger import set_trace
    #         #set_trace()
    #         #figure.canvas.renderer.clear()
    #     super(_Animation, self)._pre_draw(framedata, blit)

def draw(self):
    pass

from unittest.mock import patch

def _animate(projection, t, stars, observer, planet):
    import matplotlib.pyplot as plt

    # print(t)
    # print(t[0].utc_strftime('%Y %B'))

    o = observer.at(t)
    p = o.observe(planet)

    p0 = planet.at(t)
    mag = -8.88 + 5.0 * np.log10(p.distance().au * p0.distance().au)

    project = projection(p)
    x, y = project(p)

    fig, ax = plt.subplots()

    plt.tick_params(axis='both', which='both',
                    bottom=False, labelbottom=False,
                    left=False, labelleft=False)

    # Draw invisibly the entire path, so that the bounding boxes are
    # known ahead of time.
    ax.scatter(x, y, color='b', alpha=0.0)

    art, mag_to_radius = _plot_stars(stars, o, project, ax, 6.0, 8.5, 0.8)

    # http://aa.usno.navy.mil/software/mica/USNO-AA-TN-2003-04.pdf
    # m - M = 5 (log10 d - 1)
    # m = 5 (log10 d - 1) + M

    x_left, x_right = ax.get_xlim()
    y_bottom, y_top = ax.get_ylim()
    text_x = x_left
    text_y = y_bottom - (y_top - y_bottom) * 0.01

    saturn_color = '#f7dfae'
    #'#a69276'  # chroma('#d8c2a5').darken().hex()

    # Somehow matplotlib is (a) drawing the date text twice which ruins
    # the anti-aliasing and makes it look blocky, and (b) also never
    # erasing that part of the figure so the month names all pile up on
    # top of each other.  So let's manually paint a white rectangle
    # behind the text each time to overwrite the previous text.
    r = ax.add_patch(plt.Rectangle(
        (x_left, y_bottom - 0.012), 0.08, 0.01,
        clip_on=False, facecolor='white',
    ))

    items = []  # so update() can get the axes init() created

    def init():
        print('init()')
        planet_art = ax.scatter(x[100], y[100], color=saturn_color)
        #print(dir(planet_art))
        date_text = ax.text(text_x, text_y, '', ha='left', va='top')
        items[:] = [planet_art, date_text]
        return planet_art, r, #date_text

    def update(i):
        #print('Frame {}'.format(i))

        planet_art, date_text = items

        # planet_art.set_xdata(x[i])
        # planet_art.set_ydata(y[i])

        planet_art.set_offsets(((x[i], y[i]),))
        planet_art.set_sizes((mag_to_radius(mag[i]) ** 2.0,))

        date_text.set_text(t[i].utc_strftime('%Y %B'))

        ax.set_facecolor('black')
        return planet_art, r, date_text

    frames = len(x)
    #frames = 90

    ax.set_facecolor('black')

    anim = _Animation(fig, update, frames=frames,
                      blit=True, init_func=init,
                      interval=16)
    plt.close()
    return anim
    rcParams.validate['animation.writer'].valid['ffmpeg_small'] = 'ffmpeg_small'
    with rc_context({'animation.writer': 'ffmpeg_small'}):
        html = anim.to_html5_video()

    anim.patcher.__exit__()
    return HTML(html)

def _save_and_display(anim, path):
    Writer = _FFMpegWriter
    writer = Writer(codec='h264',
                    bitrate=rcParams['animation.bitrate'],
                    fps=1000. / anim._interval)
    anim.save(path, writer=writer)
    anim.patcher.__exit__()
    return HTML(r'''<video autoplay loop>
  <source type="video/mp4" src="{}">
  Alas! Your browser does not support the video tag.
</video>'''.format(path)) #'file://' + path))

class T(object):
    from time import time
    def __enter__(self):
        self.t0 = self.time()
    def __exit__(self, *args):
        print('<<< %s >>>' % (self.time() - self.t0))

def printout(thing):
    print()
    for line in debug(thing):
        print(line)
    print()

def debug(thing, indent=0):
    a = 'A' if thing.get_animated() else '-'
    yield '{:<{}}{} {} {}'.format('', indent, hex(id(thing)), a, repr(thing))
    indent += 2
    for child in thing.get_children():
        yield from debug(child, indent)

max_stars_repo_name: pdav-sfi/python-skyfield

max_issues_repo_name: pdav-sfi/python-skyfield

max_forks_repo_name: pdav-sfi/python-skyfield

reason: import numpy,from astropy


text:
import math
import random
import statistics
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats

def parse_roll_test(test):
	# Parse the result of 'roll test' from Minstrel Hall
	dist = {}
	total = 0
	for line in test.split("\n"):
		if ":" in line:
			val, prob = line.split(":")
			dist[val] = int(prob)
			total += int(prob)
	# Parsing complete. We now have a mapping of result to number of instances,
	# and a total result count. So dist[X]/total == probability of X occurring,
	# for any X.
	return dist, total

def chisq(dist, total):
	expected = total / len(dist)
	error = 0
	for result, count in dist.items():
		error += (count - expected) ** 2 / expected
	print("χ² =", error)

if 0: chisq(*parse_roll_test("""
1: 10135
2: 9971
3: 9774
4: 9849
5: 10059
6: 9936
7: 9990
8: 10027
9: 9917
10: 10054
11: 10202
12: 10008
13: 10136
14: 10060
15: 10012
16: 9941
17: 10007
18: 9956
19: 10096
20: 9870
"""))

def roll_dice(n):
	# Roll an N-sided dice. Mess with this to create unfair distributions.
	roll = random.randrange(n) + 1
	if roll in {1, n} and not random.randrange(40):
		roll = random.randrange(n) + 1
	return roll

def test_dice_roller(n, tries):
	dist = {}
	# Initialize to all zeroes to ensure that we have entries for everything
	for i in range(n):
		dist[i + 1] = 0
	# Roll a bunch of dice and see what comes up
	for _ in range(tries):
		dist[roll_dice(n)] += 1
	chisq(dist, tries)

# for _ in range(15): test_dice_roller(20, 200000)

#(x ** 500_000_000) * ((1-x) ** 500_000_000)

#fac(1_000_000_000) / (fac(500_000_000) * fac(500_000_000))

def choose(n, k):
	# Actually calculate n choose k
	# n! / (k!(n-k)!)
	num = denom = 1
	for i in range(1, k + 1):
		num *= i + n - k
		denom *= i
	return num // denom
def approxchoose(n, k):
	# Approximate n choose k for large n and k
	# Returns values a smidge too high, and has to be special-cased for N choose 0 and N choose N
	if k == 0 or k == n: return 1
	return (n / (2 * math.pi * k * (n-k))) ** 0.5 * (n ** n) / (k ** k * (n - k) ** (n - k))
def choosehalf(n):
	# Approximate 2n choose n
	return 2 ** (2*n) / (math.pi * n)**0.5
print(choose(20, 10))
print(approxchoose(20, 10))
print(choosehalf(10))
# print(approxchoose(1_000_000, 500_000))
# print(choosehalf(500_000))

# Binomial distribution for N coin tosses
# What is the standard deviation of [0]*N + [1]*N ?
N = 1_000_000_000
half = 5_000_000_000
def stdev(N):
	ss = 0
	for n in range(N):
		ss += approxchoose(N, n) * (n - N/2) ** 2
	return (ss / (2**N-1)) ** 0.5

def pascal(n):
	if n == 0: return [1]
	p = [0] + pascal(n - 1) + [0]
	return [p[i] + p[i + 1] for i in range(len(p) - 1)]

if 0:
	# Calculate the error in the approxchoose function by comparing against
	# the corresponding row in Pascal's Triangle
	N = 20
	for n, count in enumerate(pascal(N)):
		print(f"%{len(str(N))}d %.5f" % (n, approxchoose(N, n) / count), count)
	import sys; sys.exit(0)

# statistics.NormalDist(0.5, stdev(N))
# for N in (4, 10, 20, 100, 1_000_000, 1_000_000_000, 1_000_000_000_000):
for N in (5, 10, 15, 20, 50, 75, 100):
	sigmaN = stdev(N) / N
	print(N, sigmaN)
	x = np.linspace(0, 1.0, N + 2)
	if N < 50:
		p = [p * n / 2**N * 2 for n, p in enumerate(pascal(N) + [0])]
		plt.plot(x, p, label=f"Actual [{N}]")
	if 0: # Slow calculation of the same figure stdev gives
		samples = []
		for n, count in enumerate(pascal(N)): samples.extend(count * [n])
		# mu / N, sigma / N
		muN, sigmaN = 0.5, statistics.stdev(samples)/N
		print(N, sigmaN)
	plt.plot(x, stats.norm.pdf(x, 0.5, sigmaN), label=str(N))
# What is the first derivative of the PDF of (1e9 choose 5e8) at 0.5?
# f''(x) = 1/sqrt(2*pi)*e**(-1/2x^2) * (x^2-1)
# 499_000_000 <= x <= 501_000_000 ?? What probability?
# CDF: What is the probability that x < 499e6 ?
# CDF: What is the probability that x < 501e6 ?
# What is the spread around the mean such that CDF(x+spread) - CDF(x-spread) == 0.99?

plt.legend()
plt.show()

max_stars_repo_name: stephenangelico/shed

max_issues_repo_name: stephenangelico/shed

max_forks_repo_name: stephenangelico/shed

reason: import numpy,import scipy


text:
from typing import Collection

import numpy as np
from cv2 import contourArea
from scipy.ndimage import gaussian_filter

from glomeruli_segmentation.data_classes import Rect, Tile
from glomeruli_segmentation.extract_results import find_contours
from glomeruli_segmentation.util import combine_masks

def _threshold(image, t):
    thresholded = np.array(image, copy=True)
    thresholded[thresholded < t] = 0
    return thresholded

def _make_circle(side: int, radius: int):
    xx, yy = np.mgrid[:side, :side]
    circle = (xx - side / 2) ** 2 + (yy - side / 2) ** 2
    return np.array((circle < radius**2), dtype=float)

def test_returns_empty_list_of_contours_for_blank_tile():
    shape = (256, 256, 1)
    rect = Rect(upper_left=(0, 0), width=shape[0], height=shape[1])
    blank_image = np.zeros((shape))
    blank_tile = Tile(image=blank_image, rect=rect)
    empty_contours = find_contours(blank_tile)

    assert isinstance(empty_contours, Collection)
    assert len(empty_contours) == 0

def test_returns_one_contour_for_circular_mask():
    shape = (256, 256, 1)
    rect = Rect(upper_left=(0, 0), width=shape[0], height=shape[1])
    circular_mask = _make_circle(shape[0], radius=100)
    circular_mask_tile = Tile(image=circular_mask, rect=rect)
    contours = find_contours(circular_mask_tile)

    assert isinstance(contours, Collection)
    assert isinstance(contours[0], Collection)
    assert len(contours[0]) > 10
    assert len(contours[0][0]) == 2
    assert len(contours) == 1

def test_returns_two_contours_for_two_circular_masks():
    shape = (256, 256, 1)
    rect_left = Rect(upper_left=(0, 0), width=shape[0], height=shape[1])
    rect_right = Rect(upper_left=(shape[0], 0), width=shape[0], height=shape[1])
    circular_mask = _make_circle(shape[0], radius=100)
    tiles = [Tile(image=circular_mask, rect=rect) for rect in (rect_left, rect_right)]
    contours = find_contours(combine_masks(tiles))

    assert isinstance(contours, Collection)
    assert isinstance(contours[0], Collection)
    assert len(contours[0]) > 10
    assert len(contours[0][0]) == 2
    assert len(contours) == 2

def test_returns_one_contour_for_blurred_circular_mask():
    shape = (256, 256, 1)
    rect = Rect(upper_left=(0, 0), width=shape[0], height=shape[1])
    circular_mask = _make_circle(shape[0], radius=100)
    blurred_circular_mask = gaussian_filter(circular_mask, sigma=(10, 10))
    blurred_circular_mask_tile = Tile(image=blurred_circular_mask, rect=rect)
    contours = find_contours(blurred_circular_mask_tile)

    assert isinstance(contours, Collection)
    assert isinstance(contours[0], Collection)
    assert len(contours[0]) > 10
    assert len(contours[0][0]) == 2
    assert len(contours) == 1

def test_returns_different_contour_for_different_thresholds():
    shape = (256, 256, 1)
    rect = Rect(upper_left=(0, 0), width=shape[0], height=shape[1])
    circular_mask = _make_circle(shape[0], radius=100)
    blurred_circular_mask = gaussian_filter(circular_mask, sigma=10)
    blurred_circular_mask_tile = Tile(image=blurred_circular_mask, rect=rect)

    low_image = _threshold(blurred_circular_mask_tile.image, 0.25)
    high_image = _threshold(blurred_circular_mask_tile.image, 0.75)
    low = Tile(low_image, rect)
    high = Tile(high_image, rect)

    contours_low = find_contours(low)
    contours_high = find_contours(high)

    assert len(contours_high) == 1
    assert len(contours_low) == 1
    assert len(contours_high[0]) > 10
    assert len(contours_high[0]) != len(contours_low[0]) or not np.allclose(contours_high, contours_low)
    print(type(contours_high[0]), contours_high[0])
    assert contourArea(np.array(contours_high[0])) < contourArea(np.array(contours_low[0]))

reason: import numpy,from scipy