431 lines
16 KiB
Python
431 lines
16 KiB
Python
|
import multiprocessing
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
from functools import partial
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import tensorflow as tf
|
||
|
|
from baselines.common.tf_util import normc_initializer
|
||
|
|
from mpi4py import MPI
|
||
|
|
import tensorflow_probability as tfp
|
||
|
|
import os
|
||
|
|
import numpy as np
|
||
|
|
tfd = tfp.distributions
|
||
|
|
|
||
|
|
layers = tf.keras.layers
|
||
|
|
|
||
|
|
|
||
|
|
def bcast_tf_vars_from_root(sess, vars):
|
||
|
|
"""
|
||
|
|
Send the root node's parameters to every worker.
|
||
|
|
|
||
|
|
Arguments:
|
||
|
|
sess: the TensorFlow session.
|
||
|
|
vars: all parameter variables including optimizer's
|
||
|
|
"""
|
||
|
|
rank = MPI.COMM_WORLD.Get_rank()
|
||
|
|
for var in vars:
|
||
|
|
if rank == 0:
|
||
|
|
MPI.COMM_WORLD.bcast(sess.run(var))
|
||
|
|
else:
|
||
|
|
sess.run(tf.assign(var, MPI.COMM_WORLD.bcast(None)))
|
||
|
|
|
||
|
|
|
||
|
|
def get_mean_and_std(array):
|
||
|
|
comm = MPI.COMM_WORLD
|
||
|
|
task_id, num_tasks = comm.Get_rank(), comm.Get_size()
|
||
|
|
local_mean = np.array(np.mean(array))
|
||
|
|
sum_of_means = np.zeros((), dtype=np.float32)
|
||
|
|
comm.Allreduce(local_mean, sum_of_means, op=MPI.SUM)
|
||
|
|
mean = sum_of_means / num_tasks
|
||
|
|
|
||
|
|
n_array = array - mean
|
||
|
|
sqs = n_array ** 2
|
||
|
|
local_mean = np.array(np.mean(sqs))
|
||
|
|
sum_of_means = np.zeros((), dtype=np.float32)
|
||
|
|
comm.Allreduce(local_mean, sum_of_means, op=MPI.SUM)
|
||
|
|
var = sum_of_means / num_tasks
|
||
|
|
std = var ** 0.5
|
||
|
|
return mean, std
|
||
|
|
|
||
|
|
|
||
|
|
def guess_available_gpus(n_gpus=None):
|
||
|
|
if n_gpus is not None:
|
||
|
|
return list(range(n_gpus))
|
||
|
|
if 'CUDA_VISIBLE_DEVICES' in os.environ:
|
||
|
|
cuda_visible_divices = os.environ['CUDA_VISIBLE_DEVICES']
|
||
|
|
cuda_visible_divices = cuda_visible_divices.split(',')
|
||
|
|
return [int(n) for n in cuda_visible_divices]
|
||
|
|
nvidia_dir = '/proc/driver/nvidia/gpus/'
|
||
|
|
if os.path.exists(nvidia_dir):
|
||
|
|
n_gpus = len(os.listdir(nvidia_dir))
|
||
|
|
return list(range(n_gpus))
|
||
|
|
raise Exception("Couldn't guess the available gpus on this machine")
|
||
|
|
|
||
|
|
|
||
|
|
def setup_mpi_gpus():
|
||
|
|
"""
|
||
|
|
Set CUDA_VISIBLE_DEVICES using MPI.
|
||
|
|
"""
|
||
|
|
available_gpus = guess_available_gpus()
|
||
|
|
|
||
|
|
node_id = platform.node()
|
||
|
|
nodes_ordered_by_rank = MPI.COMM_WORLD.allgather(node_id)
|
||
|
|
processes_outranked_on_this_node = [n for n in nodes_ordered_by_rank[:MPI.COMM_WORLD.Get_rank()] if n == node_id]
|
||
|
|
local_rank = len(processes_outranked_on_this_node)
|
||
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = str(available_gpus[local_rank])
|
||
|
|
|
||
|
|
|
||
|
|
def guess_available_cpus():
|
||
|
|
return int(multiprocessing.cpu_count())
|
||
|
|
|
||
|
|
|
||
|
|
def setup_tensorflow_session():
|
||
|
|
num_cpu = guess_available_cpus()
|
||
|
|
|
||
|
|
tf_config = tf.ConfigProto(
|
||
|
|
inter_op_parallelism_threads=num_cpu,
|
||
|
|
intra_op_parallelism_threads=num_cpu
|
||
|
|
)
|
||
|
|
tf_config.gpu_options.allow_growth = True
|
||
|
|
return tf.Session(config=tf_config)
|
||
|
|
|
||
|
|
|
||
|
|
def random_agent_ob_mean_std(env, nsteps=10000):
|
||
|
|
ob = np.asarray(env.reset())
|
||
|
|
if MPI.COMM_WORLD.Get_rank() == 0:
|
||
|
|
obs = [ob]
|
||
|
|
for _ in range(nsteps):
|
||
|
|
ac = env.action_space.sample()
|
||
|
|
ob, _, done, _ = env.step(ac)
|
||
|
|
if done:
|
||
|
|
ob = env.reset()
|
||
|
|
obs.append(np.asarray(ob))
|
||
|
|
mean = np.mean(obs, 0).astype(np.float32)
|
||
|
|
std = np.std(obs, 0).mean().astype(np.float32)
|
||
|
|
else:
|
||
|
|
mean = np.empty(shape=ob.shape, dtype=np.float32)
|
||
|
|
std = np.empty(shape=(), dtype=np.float32)
|
||
|
|
MPI.COMM_WORLD.Bcast(mean, root=0)
|
||
|
|
MPI.COMM_WORLD.Bcast(std, root=0)
|
||
|
|
return mean, std
|
||
|
|
|
||
|
|
|
||
|
|
def layernorm(x):
|
||
|
|
m, v = tf.nn.moments(x, -1, keep_dims=True)
|
||
|
|
return (x - m) / (tf.sqrt(v) + 1e-8)
|
||
|
|
|
||
|
|
|
||
|
|
getsess = tf.get_default_session
|
||
|
|
|
||
|
|
fc = partial(tf.layers.dense, kernel_initializer=normc_initializer(1.))
|
||
|
|
activ = tf.nn.relu
|
||
|
|
|
||
|
|
|
||
|
|
def flatten_two_dims(x):
|
||
|
|
return tf.reshape(x, [-1] + x.get_shape().as_list()[2:])
|
||
|
|
|
||
|
|
|
||
|
|
def unflatten_first_dim(x, sh):
|
||
|
|
return tf.reshape(x, [sh[0], sh[1]] + x.get_shape().as_list()[1:])
|
||
|
|
|
||
|
|
|
||
|
|
def add_pos_bias(x):
|
||
|
|
with tf.variable_scope(name_or_scope=None, default_name="pos_bias"):
|
||
|
|
b = tf.get_variable(name="pos_bias", shape=[1] + x.get_shape().as_list()[1:], dtype=tf.float32,
|
||
|
|
initializer=tf.zeros_initializer())
|
||
|
|
return x + b
|
||
|
|
|
||
|
|
|
||
|
|
def small_convnet(x, nl, feat_dim, last_nl, layernormalize, batchnorm=False):
|
||
|
|
# nl=512, feat_dim=None, last_nl=0, layernormalize=0, batchnorm=False
|
||
|
|
bn = tf.layers.batch_normalization if batchnorm else lambda x: x
|
||
|
|
x = bn(tf.layers.conv2d(x, filters=32, kernel_size=8, strides=(4, 4), activation=nl))
|
||
|
|
x = bn(tf.layers.conv2d(x, filters=64, kernel_size=4, strides=(2, 2), activation=nl))
|
||
|
|
x = bn(tf.layers.conv2d(x, filters=64, kernel_size=3, strides=(1, 1), activation=nl))
|
||
|
|
x = tf.reshape(x, (-1, np.prod(x.get_shape().as_list()[1:])))
|
||
|
|
x = bn(fc(x, units=feat_dim, activation=None))
|
||
|
|
if last_nl is not None:
|
||
|
|
x = last_nl(x)
|
||
|
|
if layernormalize:
|
||
|
|
x = layernorm(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
# new add
|
||
|
|
class SmallConv(tf.keras.Model):
|
||
|
|
def __init__(self, feat_dim, name=None):
|
||
|
|
super(SmallConv, self).__init__(name=name)
|
||
|
|
self.conv1 = layers.Conv2D(filters=32, kernel_size=8, strides=(4, 4), activation=tf.nn.leaky_relu)
|
||
|
|
self.conv2 = layers.Conv2D(filters=64, kernel_size=4, strides=(2, 2), activation=tf.nn.leaky_relu)
|
||
|
|
self.conv3 = layers.Conv2D(filters=64, kernel_size=3, strides=(1, 1), activation=tf.nn.leaky_relu)
|
||
|
|
self.fc = layers.Dense(units=feat_dim, activation=None)
|
||
|
|
|
||
|
|
def call(self, x):
|
||
|
|
x = self.conv1(x)
|
||
|
|
x = self.conv2(x)
|
||
|
|
x = self.conv3(x)
|
||
|
|
x = tf.reshape(x, (-1, np.prod(x.get_shape().as_list()[1:])))
|
||
|
|
x = self.fc(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
# new add
|
||
|
|
class ResBlock(tf.keras.Model):
|
||
|
|
def __init__(self, hidsize):
|
||
|
|
super(ResBlock, self).__init__()
|
||
|
|
self.hidsize = hidsize
|
||
|
|
self.dense1 = layers.Dense(hidsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.dense2 = layers.Dense(hidsize, activation=None)
|
||
|
|
|
||
|
|
def call(self, xs):
|
||
|
|
x, a = xs
|
||
|
|
res = self.dense1(tf.concat([x, a], axis=-1))
|
||
|
|
res = self.dense2(tf.concat([res, a], axis=-1))
|
||
|
|
assert x.get_shape().as_list()[-1] == self.hidsize and res.get_shape().as_list()[-1] == self.hidsize
|
||
|
|
return x + res
|
||
|
|
|
||
|
|
|
||
|
|
# new add
|
||
|
|
class TransitionNetwork(tf.keras.Model):
|
||
|
|
def __init__(self, hidsize=256, name=None):
|
||
|
|
super(TransitionNetwork, self).__init__(name=name)
|
||
|
|
self.hidsize = hidsize
|
||
|
|
self.dense1 = layers.Dense(hidsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.residual_block1 = ResBlock(hidsize)
|
||
|
|
self.residual_block2 = ResBlock(hidsize)
|
||
|
|
self.dense2 = layers.Dense(hidsize, activation=None)
|
||
|
|
|
||
|
|
def call(self, xs):
|
||
|
|
s, a = xs
|
||
|
|
sh = tf.shape(a) # sh=(None,None,4)
|
||
|
|
assert len(s.get_shape().as_list()) == 3 and s.get_shape().as_list()[-1] in [512, 256]
|
||
|
|
assert len(a.get_shape().as_list()) == 3
|
||
|
|
|
||
|
|
x = flatten_two_dims(s) # shape=(None,512)
|
||
|
|
a = flatten_two_dims(a) # shape=(None,4)
|
||
|
|
|
||
|
|
#
|
||
|
|
x = self.dense1(tf.concat([x, a], axis=-1)) # (None, 256)
|
||
|
|
x = self.residual_block1([x, a]) # (None, 256)
|
||
|
|
x = self.residual_block2([x, a]) # (None, 256)
|
||
|
|
x = self.dense2(tf.concat([x, a], axis=-1)) # (None, 256)
|
||
|
|
x = unflatten_first_dim(x, sh) # shape=(None, None, 256)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class GenerativeNetworkGaussianFix(tf.keras.Model):
|
||
|
|
def __init__(self, hidsize=256, outsize=512, name=None):
|
||
|
|
super(GenerativeNetworkGaussianFix, self).__init__(name=name)
|
||
|
|
self.outsize = outsize
|
||
|
|
self.dense1 = layers.Dense(hidsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.dense2 = layers.Dense(outsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.var_single = tf.Variable(1.0, trainable=True)
|
||
|
|
|
||
|
|
self.residual_block1 = tf.keras.Sequential([
|
||
|
|
layers.Dense(hidsize, activation=tf.nn.leaky_relu), # 256
|
||
|
|
layers.Dense(hidsize, activation=None)
|
||
|
|
])
|
||
|
|
self.residual_block2 = tf.keras.Sequential([
|
||
|
|
layers.Dense(hidsize, activation=tf.nn.leaky_relu), # 256
|
||
|
|
layers.Dense(hidsize, activation=None)
|
||
|
|
])
|
||
|
|
self.residual_block3 = tf.keras.Sequential([
|
||
|
|
layers.Dense(outsize, activation=tf.nn.leaky_relu), # 512
|
||
|
|
layers.Dense(outsize, activation=None)
|
||
|
|
])
|
||
|
|
|
||
|
|
def call(self, z):
|
||
|
|
sh = tf.shape(z) # z, sh=(None,None,128)
|
||
|
|
assert z.get_shape().as_list()[-1] == 128 and len(z.get_shape().as_list()) == 3
|
||
|
|
z = flatten_two_dims(z) # shape=(None,128)
|
||
|
|
|
||
|
|
x = self.dense1(z) # (None, 256)
|
||
|
|
x = x + self.residual_block1(x) # (None, 256)
|
||
|
|
x = x + self.residual_block2(x) # (None, 256)
|
||
|
|
|
||
|
|
# variance
|
||
|
|
var_tile = tf.tile(tf.expand_dims(tf.expand_dims(self.var_single, axis=0), axis=0), [16*128, self.outsize])
|
||
|
|
|
||
|
|
# mean
|
||
|
|
x = self.dense2(x) # (None, 512)
|
||
|
|
x = x + self.residual_block3(x) # (None, 512) mean
|
||
|
|
|
||
|
|
# concat and return
|
||
|
|
x = tf.concat([x, var_tile], axis=-1) # (None, 1024)
|
||
|
|
x = unflatten_first_dim(x, sh) # shape=(None, None, 1024)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class GenerativeNetworkGaussian(tf.keras.Model):
|
||
|
|
def __init__(self, hidsize=256, outsize=512, name=None):
|
||
|
|
super(GenerativeNetworkGaussian, self).__init__(name=name)
|
||
|
|
self.dense1 = layers.Dense(hidsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.dense2 = layers.Dense(outsize, activation=tf.nn.leaky_relu)
|
||
|
|
self.dense3 = layers.Dense(outsize*2, activation=tf.nn.leaky_relu)
|
||
|
|
|
||
|
|
self.residual_block1 = tf.keras.Sequential([
|
||
|
|
layers.Dense(hidsize, activation=tf.nn.leaky_relu), # 256
|
||
|
|
layers.Dense(hidsize, activation=None)
|
||
|
|
])
|
||
|
|
self.residual_block2 = tf.keras.Sequential([
|
||
|
|
layers.Dense(hidsize, activation=tf.nn.leaky_relu), # 256
|
||
|
|
layers.Dense(hidsize, activation=None)
|
||
|
|
])
|
||
|
|
self.residual_block3 = tf.keras.Sequential([
|
||
|
|
layers.Dense(outsize, activation=tf.nn.leaky_relu), # 512
|
||
|
|
layers.Dense(outsize, activation=None)
|
||
|
|
])
|
||
|
|
|
||
|
|
def call(self, z):
|
||
|
|
sh = tf.shape(z) # z, sh=(None,None,128)
|
||
|
|
assert z.get_shape().as_list()[-1] == 128 and len(z.get_shape().as_list()) == 3
|
||
|
|
z = flatten_two_dims(z) # shape=(None,128)
|
||
|
|
|
||
|
|
x = self.dense1(z) # (None, 256)
|
||
|
|
x = x + self.residual_block1(x) # (None, 256)
|
||
|
|
x = x + self.residual_block2(x) # (None, 256)
|
||
|
|
x = self.dense2(x) # (None, 512)
|
||
|
|
x = x + self.residual_block3(x) # (None, 512)
|
||
|
|
x = self.dense3(x) # (None, 1024)
|
||
|
|
x = unflatten_first_dim(x, sh) # shape=(None, None, 1024)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class ProjectionHead(tf.keras.Model):
|
||
|
|
def __init__(self, name=None):
|
||
|
|
super(ProjectionHead, self).__init__(name=name)
|
||
|
|
self.dense1 = layers.Dense(256, activation=None)
|
||
|
|
self.dense2 = layers.Dense(128, activation=None)
|
||
|
|
self.ln1 = layers.LayerNormalization()
|
||
|
|
self.ln2 = layers.LayerNormalization()
|
||
|
|
|
||
|
|
def call(self, x, ln=False):
|
||
|
|
assert x.get_shape().as_list()[-1] == 512 and len(x.get_shape().as_list()) == 3
|
||
|
|
x = flatten_two_dims(x) # shape=(None,512)
|
||
|
|
x = self.dense1(x) # shape=(None,256)
|
||
|
|
x = self.ln1(x) # layer norm
|
||
|
|
x = tf.nn.relu(x) # relu
|
||
|
|
x = self.dense2(x) # shape=(None,128)
|
||
|
|
x = self.ln2(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class ContrastiveHead(tf.keras.Model):
|
||
|
|
def __init__(self, temperature, z_dim=128, name=None):
|
||
|
|
super(ContrastiveHead, self).__init__(name=name)
|
||
|
|
self.W = tf.Variable(tf.random.uniform((z_dim, z_dim)), name='W_Contras')
|
||
|
|
self.temperature = temperature
|
||
|
|
|
||
|
|
def call(self, z_a_pos):
|
||
|
|
z_a, z_pos = z_a_pos
|
||
|
|
Wz = tf.linalg.matmul(self.W, z_pos, transpose_b=True) # (z_dim,B) Wz.shape = (50,32)
|
||
|
|
logits = tf.linalg.matmul(z_a, Wz) # (B,B) logits.shape = (32,32)
|
||
|
|
logits = logits - tf.reduce_max(logits, 1)[:, None] # logits
|
||
|
|
logits = logits * self.temperature
|
||
|
|
return logits
|
||
|
|
|
||
|
|
|
||
|
|
def rec_log_prob(rec_params, s_next, min_sigma=1e-2):
|
||
|
|
# rec_params.shape = (None, None, 1024)
|
||
|
|
distr = normal_parse_params(rec_params, min_sigma)
|
||
|
|
log_prob = distr.log_prob(s_next) # (None, None, 512)
|
||
|
|
assert len(log_prob.get_shape().as_list()) == 3 and log_prob.get_shape().as_list()[-1] == 512
|
||
|
|
return tf.reduce_sum(log_prob, axis=-1)
|
||
|
|
|
||
|
|
|
||
|
|
def normal_parse_params(params, min_sigma=0.0):
|
||
|
|
n = params.shape[0]
|
||
|
|
d = params.shape[-1] # channel
|
||
|
|
mu = params[..., :d // 2] #
|
||
|
|
sigma_params = params[..., d // 2:]
|
||
|
|
sigma = tf.math.softplus(sigma_params)
|
||
|
|
sigma = tf.clip_by_value(t=sigma, clip_value_min=min_sigma, clip_value_max=1e5)
|
||
|
|
|
||
|
|
distr = tfd.Normal(loc=mu, scale=sigma) #
|
||
|
|
return distr
|
||
|
|
|
||
|
|
|
||
|
|
def tile_images(array, n_cols=None, max_images=None, div=1):
|
||
|
|
if max_images is not None:
|
||
|
|
array = array[:max_images]
|
||
|
|
if len(array.shape) == 4 and array.shape[3] == 1:
|
||
|
|
array = array[:, :, :, 0]
|
||
|
|
assert len(array.shape) in [3, 4], "wrong number of dimensions - shape {}".format(array.shape)
|
||
|
|
if len(array.shape) == 4:
|
||
|
|
assert array.shape[3] == 3, "wrong number of channels- shape {}".format(array.shape)
|
||
|
|
if n_cols is None:
|
||
|
|
n_cols = max(int(np.sqrt(array.shape[0])) // div * div, div)
|
||
|
|
n_rows = int(np.ceil(float(array.shape[0]) / n_cols))
|
||
|
|
|
||
|
|
def cell(i, j):
|
||
|
|
ind = i * n_cols + j
|
||
|
|
return array[ind] if ind < array.shape[0] else np.zeros(array[0].shape)
|
||
|
|
|
||
|
|
def row(i):
|
||
|
|
return np.concatenate([cell(i, j) for j in range(n_cols)], axis=1)
|
||
|
|
|
||
|
|
return np.concatenate([row(i) for i in range(n_rows)], axis=0)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
import distutils.spawn
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
|
||
|
|
def save_np_as_mp4(frames, filename, frames_per_sec=30):
|
||
|
|
print(filename)
|
||
|
|
if distutils.spawn.find_executable('avconv') is not None:
|
||
|
|
backend = 'avconv'
|
||
|
|
elif distutils.spawn.find_executable('ffmpeg') is not None:
|
||
|
|
backend = 'ffmpeg'
|
||
|
|
else:
|
||
|
|
raise NotImplementedError(
|
||
|
|
"""Found neither the ffmpeg nor avconv executables. On OS X, you can install ffmpeg via `brew install ffmpeg`. On most Ubuntu variants, `sudo apt-get install ffmpeg` should do it. On Ubuntu 14.04, however, you'll need to install avconv with `sudo apt-get install libav-tools`.""")
|
||
|
|
|
||
|
|
h, w = frames[0].shape[:2]
|
||
|
|
output_path = filename
|
||
|
|
cmdline = (backend,
|
||
|
|
'-nostats',
|
||
|
|
'-loglevel', 'error', # suppress warnings
|
||
|
|
'-y',
|
||
|
|
'-r', '%d' % frames_per_sec,
|
||
|
|
|
||
|
|
# input
|
||
|
|
'-f', 'rawvideo',
|
||
|
|
'-s:v', '{}x{}'.format(w, h),
|
||
|
|
'-pix_fmt', 'rgb24',
|
||
|
|
'-i', '-', # this used to be /dev/stdin, which is not Windows-friendly
|
||
|
|
|
||
|
|
# output
|
||
|
|
'-vcodec', 'libx264',
|
||
|
|
'-pix_fmt', 'yuv420p',
|
||
|
|
output_path)
|
||
|
|
|
||
|
|
print('saving ', output_path)
|
||
|
|
if hasattr(os, 'setsid'): # setsid not present on Windows
|
||
|
|
process = subprocess.Popen(cmdline, stdin=subprocess.PIPE, preexec_fn=os.setsid)
|
||
|
|
else:
|
||
|
|
process = subprocess.Popen(cmdline, stdin=subprocess.PIPE)
|
||
|
|
process.stdin.write(np.array(frames).tobytes())
|
||
|
|
process.stdin.close()
|
||
|
|
ret = process.wait()
|
||
|
|
if ret != 0:
|
||
|
|
print("VideoRecorder encoder exited with status {}".format(ret))
|
||
|
|
|
||
|
|
|
||
|
|
# ExponentialSchedule
|
||
|
|
class ExponentialSchedule(object):
|
||
|
|
def __init__(self, start_value, decay_factor, end_value, outside_value=None):
|
||
|
|
"""Exponential Schedule.
|
||
|
|
y = start_value * (1.0 - decay_factor) ^ t
|
||
|
|
"""
|
||
|
|
assert 0.0 <= decay_factor <= 1.0
|
||
|
|
self.start_value = start_value
|
||
|
|
self.decay_factor = decay_factor
|
||
|
|
self.end_value = end_value
|
||
|
|
|
||
|
|
def value(self, t):
|
||
|
|
v = self.start_value * np.power(1.0 - self.decay_factor, t/int(1e5))
|
||
|
|
return np.maximum(v, self.end_value)
|