SHARE
TWEET

Untitled

a guest May 24th, 2019 89 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import torch
  2. import torch.nn.functional as F
  3.  
  4. movement_filters = torch.Tensor([
  5.     [
  6.         [0, 1, 0],
  7.         [0, 0, 0],
  8.         [0, 0, 0],
  9.     ],
  10.     [
  11.         [0, 0, 0],
  12.         [0, 0, 1],
  13.         [0, 0, 0],
  14.     ],
  15.     [
  16.         [0, 0, 0],
  17.         [0, 0, 0],
  18.         [0, 1, 0],
  19.     ],
  20.     [
  21.         [0, 0, 0],
  22.         [1, 0, 0],
  23.         [0, 0, 0],
  24.     ],
  25. ]).unsqueeze(1).float()
  26.  
  27. heads = torch.zeros((4, 1, 5, 5))
  28. # Place heads in center of 4 5x5 environments
  29. heads[:, 0, 2, 2] = 1
  30.  
  31. # Each moves in a different cardinal direction
  32. actions_onehot = torch.zeros((4, 4))
  33. actions_onehot[torch.arange(4), torch.arange(4)] = 1
  34.  
  35. intermediate = F.conv2d(heads, movement_filters, padding=1)
  36.  
  37. heads = torch.einsum('bchw,bc->bhw', [intermediate, actions_onehot]).unsqueeze(1)
  38.  
  39. print(heads[:, 0])
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!
 
Top