evasion.py

   1 #!/usr/bin/env python
   2
   3 import torch
   4
   5 from torch.nn import functional as F
   6
   7 ######################################################################
   8
   9
  10 def generate_sequence(nb, height=6, width=6, T=10):
  11     rnd = torch.rand(nb, height, width)
  12     rnd[:, 0, :] = 0
  13     rnd[:, -1, :] = 0
  14     rnd[:, :, 0] = 0
  15     rnd[:, :, -1] = 0
  16     wall = 0
  17
  18     for k in range(3):
  19         wall = wall + (
  20             rnd.flatten(1).argmax(dim=1)[:, None]
  21             == torch.arange(rnd.flatten(1).size(1))[None, :]
  22         ).long().reshape(rnd.size())
  23         rnd = rnd * (1 - wall.clamp(max=1))
  24
  25     seq = wall[:, None, :, :].expand(-1, T, -1, -1).clone()
  26
  27     agent = torch.zeros(seq.size(), dtype=torch.int64)
  28     agent[:, 0, 0, 0] = 1
  29     agent_actions = torch.randint(5, (nb, T))
  30     monster = torch.zeros(seq.size(), dtype=torch.int64)
  31     monster[:, 0, -1, -1] = 1
  32     monster_actions = torch.randint(5, (nb, T))
  33
  34     all_moves = agent.new(nb, 5, height, width)
  35     for t in range(T - 1):
  36         all_moves.zero_()
  37         all_moves[:, 0] = agent[:, t]
  38         all_moves[:, 1, 1:, :] = agent[:, t, :-1, :]
  39         all_moves[:, 2, :-1, :] = agent[:, t, 1:, :]
  40         all_moves[:, 3, :, 1:] = agent[:, t, :, :-1]
  41         all_moves[:, 4, :, :-1] = agent[:, t, :, 1:]
  42         a = F.one_hot(agent_actions[:, t], num_classes=5)[:, :, None, None]
  43         after_move = (all_moves * a).sum(dim=1)
  44         collision = (
  45             (after_move * (1 - wall) * (1 - monster[:, t]))
  46             .flatten(1)
  47             .sum(dim=1)[:, None, None]
  48             == 0
  49         ).long()
  50         agent[:, t + 1] = collision * agent[:, t] + (1 - collision) * after_move
  51
  52         all_moves.zero_()
  53         all_moves[:, 0] = monster[:, t]
  54         all_moves[:, 1, 1:, :] = monster[:, t, :-1, :]
  55         all_moves[:, 2, :-1, :] = monster[:, t, 1:, :]
  56         all_moves[:, 3, :, 1:] = monster[:, t, :, :-1]
  57         all_moves[:, 4, :, :-1] = monster[:, t, :, 1:]
  58         a = F.one_hot(monster_actions[:, t], num_classes=5)[:, :, None, None]
  59         after_move = (all_moves * a).sum(dim=1)
  60         collision = (
  61             (after_move * (1 - wall) * (1 - agent[:, t + 1]))
  62             .flatten(1)
  63             .sum(dim=1)[:, None, None]
  64             == 0
  65         ).long()
  66         monster[:, t + 1] = collision * monster[:, t] + (1 - collision) * after_move
  67
  68     seq += 2 * agent + 3 * monster
  69
  70     return seq, agent_actions
  71
  72
  73 ######################################################################
  74
  75
  76 def seq2str(seq, actions=None):
  77     # symbols=" #@$"
  78     symbols = " █@$"
  79
  80     hline = ("+" + "-" * seq.size(-1)) * seq.size(1) + "+" + "\n"
  81
  82     result = hline
  83
  84     for n in range(seq.size(0)):
  85         for i in range(seq.size(2)):
  86             result += (
  87                 "|"
  88                 + "|".join(
  89                     ["".join([symbols[v.item()] for v in row]) for row in seq[n, :, i]]
  90                 )
  91                 + "|"
  92                 + "\n"
  93             )
  94
  95         result += hline
  96
  97         if actions is not None:
  98             result += (
  99                 "|"
 100                 + "|".join(
 101                     ["INESW"[a.item()] + " " * (seq.size(-1) - 1) for a in actions[n]]
 102                 )
 103                 + "|"
 104                 + "\n"
 105             )
 106
 107         result += hline
 108
 109     return result
 110
 111
 112 ######################################################################
 113
 114 if __name__ == "__main__":
 115     seq, actions = generate_sequence(40, 4, 6, T=20)
 116
 117     print(seq2str(seq, actions))