Image-based reinforcement learning with neural networks

Question

After seeing that "OpenAIGym" is not exactly supported on Windows and playing around with https://www.wolfram.com/language/12/neural-network-framework/train-an-agent-in-a-reinforcement-learning-environment.html, I decided to create a new gym as part of my reinforcement learning studies!

---

I was progressing smoothly with the first environment Pixels-v1:

• DeviceFramework API
• human interaction

then I got stuck with the actual reinforcement learning using neural networks o__O - training produces results that just drift off screen rather than aim for the center.

Is that perhaps because of picking images as input domain instead of a tiny vector of numbers summarizing the state - please see below for details and help me out?!

I have prepared a notebook to make you see the issue right away.

---

Notebook

https://www.wolframcloud.com/obj/137b4dd6-937a-47ad-b0b0-10af0ae06ed6

1. evaluate all initialization cells
2. scroll to the bottom of the notebook

(Please make and use your own copy if you can.)

---

Context

The first environment (Pixels-v1) is going to be that of a few happy pixels surviving against various simple hazards.

• "ObservedState": a 40*40 Image
• "ActionSpace": {Left, Right, Up, Down}
• "Step": I am not sure how to define the reward here for the neural network to converge, intention for the simplest case: Reward == 1 if stepped closer than ever before to center, otherwise 0; also Ended == True if active Pixel hit edges.

Questions

I would expect my policy network to use the environment image directly.

policyInput = NetEncoder[{"Image", {40, 40}, "Grayscale"}];
policyOutput = NetDecoder[actionClass];
policy = NetInitialize@NetChain[
   {
    4, SoftmaxLayer[]
    },
   "Input" -> policyInput,
   "Output" -> policyOutput
   ]
policyOutput;
policy // netSize

1. Should I change the "ObservedState" from being an image of binary 40*40 image with the active pixel being a single pixel to something more noticable to the neural net?
2. Should I change the policy network definition (snippet above), how?
3. Should I change the loss function definition, how?
4. Should I change the step function definition (snippet below), how, especially with regards to the "Reward"?

(where \[Rho][[i]] is the active pixel.)

---

I would strongly appreciate if someone could provide detailed guidance on what I am missing at this point. I would love to make Mathematica more accessible for beginners of reinforcement learning with neural networks.

Notebook: https://www.wolframcloud.com/obj/137b4dd6-937a-47ad-b0b0-10af0ae06ed6

1. evaluate all initialization cells and
2. scroll to the bottom of the notebook

(Please make and use your own copy if you can. Same link as above.)

---

(If this is not the right place / way to pose this question, please tell me how to correct / where to move it and I will oblige.)