Update.
[picoclvr.git] / world.py
1 #!/usr/bin/env python
2
3 import math, sys
4
5 import torch, torchvision
6
7 from torch import nn
8 from torch.nn import functional as F
9 import cairo
10
11
12 class Box:
13     def __init__(self, x, y, w, h, r, g, b):
14         self.x = x
15         self.y = y
16         self.w = w
17         self.h = h
18         self.r = r
19         self.g = g
20         self.b = b
21
22     def collision(self, scene):
23         for c in scene:
24             if (
25                 self is not c
26                 and max(self.x, c.x) <= min(self.x + self.w, c.x + c.w)
27                 and max(self.y, c.y) <= min(self.y + self.h, c.y + c.h)
28             ):
29                 return True
30         return False
31
32
33 def scene2tensor(xh, yh, scene, size):
34     width, height = size, size
35     pixel_map = torch.ByteTensor(width, height, 4).fill_(255)
36     data = pixel_map.numpy()
37     surface = cairo.ImageSurface.create_for_data(
38         data, cairo.FORMAT_ARGB32, width, height
39     )
40
41     ctx = cairo.Context(surface)
42     ctx.set_fill_rule(cairo.FILL_RULE_EVEN_ODD)
43
44     for b in scene:
45         ctx.move_to(b.x * size, b.y * size)
46         ctx.rel_line_to(b.w * size, 0)
47         ctx.rel_line_to(0, b.h * size)
48         ctx.rel_line_to(-b.w * size, 0)
49         ctx.close_path()
50         ctx.set_source_rgba(b.r, b.g, b.b, 1.0)
51         ctx.fill()
52
53     hs = size * 0.1
54     ctx.set_source_rgba(0.0, 0.0, 0.0, 1.0)
55     ctx.move_to(xh * size - hs / 2, yh * size - hs / 2)
56     ctx.rel_line_to(hs, 0)
57     ctx.rel_line_to(0, hs)
58     ctx.rel_line_to(-hs, 0)
59     ctx.close_path()
60     ctx.fill()
61
62     return pixel_map[None, :, :, :3].flip(-1).permute(0, 3, 1, 2).float() / 255
63
64
65 def random_scene():
66     scene = []
67     colors = [
68         (1.00, 0.00, 0.00),
69         (0.00, 1.00, 0.00),
70         (0.60, 0.60, 1.00),
71         (1.00, 1.00, 0.00),
72         (0.75, 0.75, 0.75),
73     ]
74
75     for k in range(10):
76         wh = torch.rand(2) * 0.2 + 0.2
77         xy = torch.rand(2) * (1 - wh)
78         c = colors[torch.randint(len(colors), (1,))]
79         b = Box(
80             xy[0].item(), xy[1].item(), wh[0].item(), wh[1].item(), c[0], c[1], c[2]
81         )
82         if not b.collision(scene):
83             scene.append(b)
84
85     return scene
86
87
88 def generate_sequence(nb_steps=10, all_frames=False, size=64):
89     delta = 0.1
90     effects = [
91         (False, 0, 0),
92         (False, delta, 0),
93         (False, 0, delta),
94         (False, -delta, 0),
95         (False, 0, -delta),
96         (True, delta, 0),
97         (True, 0, delta),
98         (True, -delta, 0),
99         (True, 0, -delta),
100     ]
101
102     while True:
103         frames = []
104
105         scene = random_scene()
106         xh, yh = tuple(x.item() for x in torch.rand(2))
107
108         frames.append(scene2tensor(xh, yh, scene, size=size))
109
110         actions = torch.randint(len(effects), (nb_steps,))
111         change = False
112
113         for a in actions:
114             g, dx, dy = effects[a]
115             if g:
116                 for b in scene:
117                     if b.x <= xh and b.x + b.w >= xh and b.y <= yh and b.y + b.h >= yh:
118                         x, y = b.x, b.y
119                         b.x += dx
120                         b.y += dy
121                         if (
122                             b.x < 0
123                             or b.y < 0
124                             or b.x + b.w > 1
125                             or b.y + b.h > 1
126                             or b.collision(scene)
127                         ):
128                             b.x, b.y = x, y
129                         else:
130                             xh += dx
131                             yh += dy
132                             change = True
133             else:
134                 x, y = xh, yh
135                 xh += dx
136                 yh += dy
137                 if xh < 0 or xh > 1 or yh < 0 or yh > 1:
138                     xh, yh = x, y
139
140             if all_frames:
141                 frames.append(scene2tensor(xh, yh, scene, size=size))
142
143         if not all_frames:
144             frames.append(scene2tensor(xh, yh, scene, size=size))
145
146         if change:
147             break
148
149     return frames, actions
150
151
152 ######################################################################
153
154
155 # ||x_i - c_j||^2 = ||x_i||^2 + ||c_j||^2 - 2<x_i, c_j>
156 def sq2matrix(x, c):
157     nx = x.pow(2).sum(1)
158     nc = c.pow(2).sum(1)
159     return nx[:, None] + nc[None, :] - 2 * x @ c.t()
160
161
162 def update_centroids(x, c, nb_min=1):
163     _, b = sq2matrix(x, c).min(1)
164     b.squeeze_()
165     nb_resets = 0
166
167     for k in range(0, c.size(0)):
168         i = b.eq(k).nonzero(as_tuple=False).squeeze()
169         if i.numel() >= nb_min:
170             c[k] = x.index_select(0, i).mean(0)
171         else:
172             n = torch.randint(x.size(0), (1,))
173             nb_resets += 1
174             c[k] = x[n]
175
176     return c, b, nb_resets
177
178
179 def kmeans(x, nb_centroids, nb_min=1):
180     if x.size(0) < nb_centroids * nb_min:
181         print("Not enough points!")
182         exit(1)
183
184     c = x[torch.randperm(x.size(0))[:nb_centroids]]
185     t = torch.full((x.size(0),), -1)
186     n = 0
187
188     while True:
189         c, u, nb_resets = update_centroids(x, c, nb_min)
190         n = n + 1
191         nb_changes = (u - t).sign().abs().sum() + nb_resets
192         t = u
193         if nb_changes == 0:
194             break
195
196     return c, t
197
198
199 ######################################################################
200
201
202 def patchify(x, factor, invert_size=None):
203     if invert_size is None:
204         return (
205             x.reshape(
206                 x.size(0),  # 0
207                 x.size(1),  # 1
208                 factor,  # 2
209                 x.size(2) // factor,  # 3
210                 factor,  # 4
211                 x.size(3) // factor,  # 5
212             )
213             .permute(0, 2, 4, 1, 3, 5)
214             .reshape(-1, x.size(1), x.size(2) // factor, x.size(3) // factor)
215         )
216     else:
217         return (
218             x.reshape(
219                 invert_size[0],  # 0
220                 factor,  # 1
221                 factor,  # 2
222                 invert_size[1],  # 3
223                 invert_size[2] // factor,  # 4
224                 invert_size[3] // factor,  # 5
225             )
226             .permute(0, 3, 1, 4, 2, 5)
227             .reshape(invert_size)
228         )
229
230
231 class Normalizer(nn.Module):
232     def __init__(self, mu, std):
233         super().__init__()
234         self.mu = nn.Parameter(mu)
235         self.log_var = nn.Parameter(2*torch.log(std))
236
237     def forward(self, x):
238         return (x-self.mu)/torch.exp(self.log_var/2.0)
239
240 class SignSTE(nn.Module):
241     def __init__(self):
242         super().__init__()
243
244     def forward(self, x):
245         # torch.sign() takes three values
246         s = (x >= 0).float() * 2 - 1
247         if self.training:
248             u = torch.tanh(x)
249             return s + u - u.detach()
250         else:
251             return s
252
253
254 def train_encoder(
255     train_input,
256     dim_hidden=64,
257     block_size=16,
258     nb_bits_per_block=10,
259     lr_start=1e-3, lr_end=1e-5,
260     nb_epochs=50,
261     batch_size=25,
262     device=torch.device("cpu"),
263 ):
264     mu, std = train_input.mean(), train_input.std()
265
266     encoder = nn.Sequential(
267         Normalizer(mu, std),
268         nn.Conv2d(3, dim_hidden, kernel_size=5, stride=1, padding=2),
269         nn.ReLU(),
270         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
271         nn.ReLU(),
272         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
273         nn.ReLU(),
274         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
275         nn.ReLU(),
276         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
277         nn.ReLU(),
278         nn.Conv2d(
279             dim_hidden,
280             nb_bits_per_block,
281             kernel_size=block_size,
282             stride=block_size,
283             padding=0,
284         ),
285         SignSTE(),
286     )
287
288     decoder = nn.Sequential(
289         nn.ConvTranspose2d(
290             nb_bits_per_block,
291             dim_hidden,
292             kernel_size=block_size,
293             stride=block_size,
294             padding=0,
295         ),
296         nn.ReLU(),
297         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
298         nn.ReLU(),
299         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
300         nn.ReLU(),
301         nn.Conv2d(dim_hidden, dim_hidden, kernel_size=5, stride=1, padding=2),
302         nn.ReLU(),
303         nn.Conv2d(dim_hidden, 3, kernel_size=5, stride=1, padding=2),
304     )
305
306     model = nn.Sequential(encoder, decoder)
307
308     nb_parameters = sum(p.numel() for p in model.parameters())
309
310     print(f"nb_parameters {nb_parameters}")
311
312     model.to(device)
313
314     for k in range(nb_epochs):
315         lr=math.exp(math.log(lr_start) + math.log(lr_end/lr_start)/(nb_epochs-1)*k)
316         print(f"lr {lr}")
317         optimizer = torch.optim.Adam(model.parameters(), lr=lr)
318         acc_loss, nb_samples = 0.0, 0
319
320         for input in train_input.split(batch_size):
321             output = model(input)
322             loss = F.mse_loss(output, input)
323             acc_loss += loss.item() * input.size(0)
324             nb_samples += input.size(0)
325
326             optimizer.zero_grad()
327             loss.backward()
328             optimizer.step()
329
330         print(f"loss {k} {acc_loss/nb_samples}")
331         sys.stdout.flush()
332
333     return encoder, decoder
334
335
336 ######################################################################
337
338 if __name__ == "__main__":
339     import time
340
341     all_frames = []
342     nb = 25000
343     start_time = time.perf_counter()
344     for n in range(nb):
345         frames, actions = generate_sequence(nb_steps=31)
346         all_frames += frames
347     end_time = time.perf_counter()
348     print(f"{nb / (end_time - start_time):.02f} samples per second")
349
350     input = torch.cat(all_frames, 0)
351     encoder, decoder = train_encoder(input)
352
353     # x = patchify(input, 8)
354     # y = x.reshape(x.size(0), -1)
355     # print(f"{x.size()=} {y.size()=}")
356     # centroids, t = kmeans(y, 4096)
357     # results = centroids[t]
358     # results = results.reshape(x.size())
359     # results = patchify(results, 8, input.size())
360
361     z = encoder(input)
362     results = decoder(z)
363
364     print(f"{input.size()=} {z.size()=} {results.size()=}")
365
366     torchvision.utils.save_image(input[:64], "orig.png", nrow=8)
367
368     torchvision.utils.save_image(results[:64], "qtiz.png", nrow=8)
369
370     # frames, actions = generate_sequence(nb_steps=31, all_frames=True)
371     # frames = torch.cat(frames, 0)
372     # torchvision.utils.save_image(frames, "seq.png", nrow=8)