def test_flowfield():
u = UniImageViewer()
x = bad_monkey()
theta = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]).expand(1, -1, -1)
grid = F.affine_grid(theta, x.shape)
out = F.grid_sample(x, grid)
u.render(out[0], block=True)
def evaluate(args, env, policy, render=False):
with torch.no_grad():
def get_action(s, prepro, transform, view, policy, action_map, device):
s = prepro(s)
s_t = transform(s).unsqueeze(0).type(policy.dtype).to(device)
kp = view(s_t)
p = softmax(keypoints.models.knn.flatten().matmul(policy), dim=0)
a = Categorical(p).sample()
a = action_map(a)
return a, kp
v = UniImageViewer()
datapack = ds.datasets[args.dataset]
if args.model_type != 'nop':
transporter_net = transporter.make(args, map_device='cpu')
view = Keypoints(transporter_net).to(args.device)
else:
view = nop
s = env.reset()
a, kp = get_action(s, datapack.prepro, datapack.transforms, view,
policy, datapack.action_map, args.device)
done = False
reward = 0.0
while not done:
s, r, done, i = env.step(a)
reward += r
a, kp = get_action(s, datapack.prepro, datapack.transforms, view,
policy, datapack.action_map, args.device)
if render:
if args.model_keypoints:
s = datapack.prepro(s)
s = TVF.to_tensor(s).unsqueeze(0)
s = plot_keypoints_on_image(kp[0], s[0])
v.render(s)
else:
env.render()
return reward
def test_tps_random():
images = []
u = UniImageViewer(screen_resolution=(2400, 1200))
x = bad_monkey()
for i in range(5, 10):
set = []
for _ in range(8):
pass
# set.append(tps_random(x, num_control_points=20, var=1 / i))
st = torch.cat(set, dim=2)
images.append(st)
img = torch.cat(images, dim=3)
u.render(img, block=True)
def test_tps():
u = UniImageViewer()
x = bad_monkey()
theta = torch.tensor([[[0.0, 0.0], [0., 0.], [0., 0.], [0., 0.], [0., 0.],
[0., 0.], [0.0, 0.0]]])
c = torch.tensor([
[0., 0],
[1., 0],
[1., 1],
[0, 1],
]).unsqueeze(0)
grid = tps_grid(theta, c, x.shape)
out = F.grid_sample(x, grid)
u.render(out[0], block=True)
def test_pong_fill():
l = 6000
display = True
ds = AtariDataset('Pong-v0', l, d.pong_prepro,
transforms = T.Compose([T.ToTensor(), T.ToPILImage()]),
end_trajectory=if_done_or_nonzero_reward)
assert len(ds) >= l
print(len(ds))
disp = UniImageViewer(screen_resolution=(512, 512))
for img1, img2 in ds:
#i = np.concatenate((img[0], img[1]), axis=1)
if display:
disp.render(img1)
time.sleep(0.03)
else:
plt.imshow(img1, cmap='gray', vmin=0, vmax=256.0)
plt.show()
def test_co_ords():
height, width = 16, 16
hm = torch.zeros(1, 1, height, width)
hm[0, 0, 0, 15] = 20.0
k, p = MF.spacial_softmax(hm, probs=True)
g = MF.gaussian_like_function(k, height, width)
#plot_heightmap3d(hm[0, 0].detach().numpy())
#plot_heightmap3d(g[0, 0].detach().numpy(), k[0, 0])
#plot_single_channel(hm[0, 0])
#plot_single_channel(g[0, 0])
d = UniImageViewer()
def test_pong():
v = UniImageViewer()
l = UniImageViewer(title='processed', screen_resolution=(32, 32))
env = gym.make('Pong-v0')
s = env.reset()
done = False
while not done:
s, r, done, info = env.step(cma_es.sample())
v.render(s)
s = d.pong_color_prepro(s)
#s = cv2.cvtColor(s, cv2.COLOR_RGB2GRAY)
#s = s[34:168, :]
#s = cv2.resize(s, dsize=(32, 32), interpolation=cv2.INTER_AREA)
l.render(s)
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from colorama import Fore, Style
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau
import torch.nn as nn
import statistics as stats
from keypoints.models import vgg, knn, autoencoder
from utils import get_lr, UniImageViewer, make_grid
from keypoints.ds import datasets as ds
from apex import amp
from config import config
scale = 4
view_in = UniImageViewer('in',
screen_resolution=(128 * 2 * scale, 128 * scale))
view_z = UniImageViewer('z',
screen_resolution=(128 // 2 * 5 * scale,
128 // 2 * 4 * scale))
def log(phase):
writer.add_scalar(f'{phase}_loss', loss.item(), global_step)
if i % args.display_freq == 0:
recon = torch.cat((x[0], x_[0]), dim=2)
latent = make_grid(z[0].unsqueeze(1), 4, 4)
if args.display:
view_in.render(recon)
view_z.render(latent)
writer.add_image(f'{phase}_recon', recon, global_step)
import cma_es
from utils import UniImageViewer, plot_keypoints_on_image
from keypoints.ds import datasets as ds
from keypoints.models import transporter, functional as KF
import torch
import config
import time
from torchvision.transforms import functional as TVF
if __name__ == '__main__':
args = config.config()
with torch.no_grad():
v = UniImageViewer()
datapack = ds.datasets[args.dataset]
transporter_net = transporter.make(args).to(args.device)
if args.load is not None:
transporter_net.load(args.load)
env = gym.make(datapack.env)
while True:
s = env.reset()
done = False
while not done:
s, r, done, i = env.step(cma_es.sample())
import pyrr import gym import pygame import cma_es import main from keypoints.models import transporter import config from utils import UniImageViewer import torch from pyrr import matrix44, Vector3 from math import floor from time import sleep from PIL import Image viewer = UniImageViewer() pygame.init() vertex_src = """ # version 330 layout(location = 0) in vec3 a_position; layout(location = 1) in vec2 a_texture; uniform mat4 model; // combined translation and rotation uniform mat4 projection; out vec3 v_color; out vec2 v_texture;
def test_dual_tps_random_batched():
u = UniImageViewer(screen_resolution=(2400, 1200))
x = bad_monkey(2)