gym の render の画像を 正しくpytorch のモデルに打っ込むヤツ

gym_rgb_image_to_torch_input.py

import torch
import numpy as np
import matplotlib.pyplot as plt

# Assuming some environment that render image...
im = env.render(mode="rgb_array",
                height=84,
                width=84,
                camera_id=0) / 255.
im = torch.tensor(im.astype(np.float32))
im = im.view(-1, 84, 84, 3).permute(0, 3, 1, 2)
im = im.detach().numpy()

# visualization
im = np.hstack([im[0, 0],
                im[0, 1],
                im[0, 2]])
plt.figure()
plt.imshow(im, cmap="gray")
plt.show()

You will get this

from this rgb image