在Numpy或Pytorch中的网格上应用坐标蒙版

如何解决在Numpy或Pytorch中的网格上应用坐标蒙版

作为强化学习问题的一部分，我需要保存GPU内存。

网格的内存大小是平方的（以像素数为单位）。但是，只有代理所在的像素很重要。通过使用这些网格的“稀疏”表示，我们节省了大量内存。

为了保存此内存，我创建了另一种矩阵格式。

我有一个面具M形[Batch，N，2]：

B：批量大小
N：代理人数
2：代理的x和y坐标。

我也有一个形状动作网格[Batch，C，H，W]：

B：批量大小
C：行动次数
H：网格的高度
W：网格宽度

目标是遮罩此动作网格，以仅获取遮罩坐标处的动作，以获得形状为[B，N，C]的遮罩网格。

以下是我想要获得的示例：

>>> import numpy as np
>>> # Example of dimension
>>> BATCH_SIZE = 2
>>> N_AGENT = 3
>>> NB_ACTION = 2
>>> H_GRID = 3
>>> W_GRID = 3

>>> action_grid_batch.shape # [BATCH_SIZE,NB_ACTION,H_GRID,W_GRID]
(2,2,3,3)

>>> action_grid_batch
array( [[[[0.4000,0.5000,0.7000],# Probability  of the action 1 on the action_grid 1 in the batch
          [0.3000,0.2000,0.1000],[0.9000,0.8000,0.7000]],[[0.6000,0.3000],# Probability  of the action 2 on the action_grid 1 in the batch
          [0.7000,0.9000],[0.1000,0.3000]]],[[[0.3000,# Probability  of the action 1 on the action_grid 2 in the batch
          [0.6000,0.7000,0.4000],0.1000]],[[0.7000,# Probability  of the action 2 on element 2 in the batch
          [0.4000,0.3000,0.6000],0.9000]]]])

>>> batch_mask_agent_position
array( [[[0,1],# Position (H,W) of the agent 1 on element 1 in the batch
         [1,W) of the agent 2 on element 1 in the batch
         [2,0]],W) of the agent 3 on element 1 in the batch 

        [[1,W) of the agent 1 on element 2 in the batch
         [1,2],W) of the agent 2 on element 2 in the batch
         [2,2]]]) # Position (H,W) of the agent 3 on element 2 in the batch 

>>> output = apply_mask_on_grid(action_grid_batch,batch_mask_agent_position) # the function I would like

>>> output.shape
(2,2) # [Batch_size,N_AGENT,N_ACTION]
>>> output
array( [[[0.5000,0.5000],# Probability of the action 1 and 2 for the agent position 1 on element 1 in the batch
         [0.2000,0.8000],# Probability of the action 1 and 2 for the agent position 2 on element 1 in the batch
         [0.9000,# Probability of the action 1 and 2 for the agent position 3 on element 1 in the batch

        [[0.7000,# Probability of the action 1 and 2 for the agent position 1 on element 2 in the batch
         [0.4000,# Probability of the action 1 and 2 for the agent position 2 on element 2 in the batch
         [0.1000,0.9000]]]) # Probability of the action 1 and 2 for the agent position 3 on element 2 in the batch

（我将在Pytorch中重新实现建议的解决方案。）

预先感谢您！