Deep RL reward design for neuron centerline extraction task

Question

As part of a bigger scope project, I'm training a RL agent that attempts to reconstruct, pixel by pixel, the trajectory of a neuron on a segmented image. To give a better insight on the task that I'm trying to achieve, here is an example of an attempt on simple vessel ground truth image (red : vessel/ground truth, green : agent trajectory, yellow : overlap between the two) :

Details on the environment and the model :

At each step, the agent is picking an action based on the output of a Deep Q-Network and chooses the next pixel to move the cursor to. It can move in any of the 8 directions (in diagonal too). The goal is to match the red curve with the green one by updating each step the position of the cursor. At the moment, I'm training the agent with a few different such ground truth images and for a fair amount of episodes and it's really struggling at sticking close to the red line. I set a maximum number of pixels to stop the agent after a while in case it doesn't approximately match the red line (which it does not). I'm used to applying RL to games like Atari etc, but here it's a bit more abstract to determine what would be a smart way to design the reward system.

Question :

At the moment, I'm using the completeness score that compares what fraction of the pixels of the dilated red line are covered by the agent's path as part of my reward system. If the completeness increases as compared to the previous step, I give a reward of +1, if it doesn't, I give a reward of -0.1 (to penalize the agent adding useless pixels). I've tried a few other things like directly using the completeness score as the reward, but this technique is the overall best I can achieve, but I believe that my system is not very optimal which might cause the poor performance, as I think that the design of the reward system is essential to the good performing of a RL agent. I'm therefore open to any suggestions on how to improve or re-think the reward system by anyone that's got a bit more experience than me in designing RL agents when it's not as obvious as for some simple games.

Thanks in advance for your help !

Answer 1

Have you tried to base the rewards on the distance to the correct trajectory? For example, when the agent deviates a lot from the desired trajectory as it does in the picture you gave, it should learn that this is unwanted behaviour. In a pixel world, you could use the Manhattan distance to compute this reward.

What are the observations provided to the agent? Does it see the trajectory it should learn to follow?

You can also experiment with various functions to shape the reward. You want to do this because in many cases a binary reward doesn't provide enough information to the agent about what it did right or wrong. I have had great success by using an exponential reward which increases as the agent gets closer to the goal or target trajectory (see this paper on reward shaping)

Furthermore, you also want to prevent the agent from learning a positive feedback loop where the agent never attempts to reach the goal. Essentially, the positive feedback loop is one where the agent learns that it can gain more reward by taking a series of action that does not reach the goal. One way to prevent this is by giving a larger reward when the agent reaches the goal. In your case, if I understand correctly, the goal is not the end of the trajectory but reaching some satisfiable % level of completeness.