Implement Deep Q-Network

[seungjaeryanlee]

This is a mostly finished implementation of Deep Q-Network on the CartPole environment that needs some improvements. Following the original paper (Mnih et al., 2015), the code has both experience replay and target network implemented.

I have identified three places of improvements:

• Parallelize training
  • The Q-network is given a batch of states and outputs a batch of Q-values for all actions (batch size x # actions). This should be transformed to a 1D vector (batch size x 1) that selects the Q-value of the action that was actually selected. In Python, tf.gather_nd seems to be the best option, but it does not seem to exist in _Raw. Is there a way to use this function?
  • _Raw.gatherNd exists
• Refactor code
  • There are a lot of commented out print() statements. These are left out to aid me as I work through the improvements, but they should be removed once I am finished.
• Remove uses of _Raw
  • As per @BradLarson 's suggestion, most _Raw commands can replaced by S4TF functions. Replace them accordingly.
• Use Huber loss
  • Huber loss is mentioned in the original paper and can help with training stability.
• Use :on for Tensors to support X10 devices
• Improve reproducibility
  • The hyperparameters were chosen via a rough random search, and the results vary heavily.
• Separate gatherNd into a separate file
  • For a more streamlined improvement to swift-api, separate the gatherNd redefinition to a new file.

[seungjaeryanlee]

I have tried various hyperparameter settings from existing DQN implementations, including eaplatanios/swift-rl, higgsfield/RL-Adventure, and TF-Agents. None of them consistently got a score of 200, so I improved the model with three additional algorithms:

• Soft Target Network Update from Lillicrap et al. (2015)
• Double DQN from Hasselt, Guez, and Silver (2015)
• Combined Experience Replay from Zhang and Sutton (2017)

With these three improvements and the TF-Agents hyperparameters, I am now getting consistent 200s!

[BradLarson]

I just pushed a commit that does a few things:

• Fixes segmentation faults on GPU eager mode execution.
• Creates an extension on Tensor for gatherNd (which I'm calling dimensionalGathering until we can figure out a better name or better way of expressing the difference in parameter names).
• Splits out the types from main.swift into a few files to better organize the components here.
• Replaces a few Numpy uses with native (and faster) Swift equivalents.
• Reformats the files per swift-format's style that we use for this repository.

The segmentation fault I was seeing was due to trying to set a Tensor at an index within a larger Tensor during append() in the ReplayBuffer. I've changed this to a simpler (and probably more performant) approach using arrays of tensors that are stacked when needed during the batch extraction. For a replay buffer, we might want to consider using pure scalars there, because they'll be faster to work with when interacting with environments that take and provide scalars. I've done this with other environments and gotten very fast speeds by converting them to Tensors only at the point where you want to use the accelerator (model training and agent logic).

There are a few places where Numpy functions are being used when there are Swift-native equivalents, and that can impact performance on tight loops around an environment. I've replaced a few of these, but I think that all the rest should also be able to be replaced with Swift functions.

This isn't always solving the environment (it seems to solve ~half of the time), so if that's unexpected, you may want to check over my replay buffer logic and other changes to make sure I didn't introduce a subtle error somewhere.

For organization, I thought there was a lot of code in main.swift, so I broke it out into separate files for the core types. I think that makes it a little easier to follow what's going on here with the various pieces.

[seungjaeryanlee]

Thank you for the wonderful commit @BradLarson !

1. I agree that the replay buffer should store numbers as Swift native types instead of Tensors and only convert them in the sample() function. I forgot Swift was fast! I will make that change.
2. Converting NumPy functions to Swift is definitely a good idea. I could not find some of the functions that I wanted which was why I have been using NumPy, but using Swift functions makes much more sense.
3. The division of code into Gathering / Agent / ReplayBuffer / main makes sense to me. I am considering moving the eval() function to the agent later.
4. The issue with the degraded performance stems from the ReplayBuffer being changed from class to struct. As far as I know, that means replayBuffer is passed by value, so the agent.replayBuffer is never updated and is empty the entire time, and thus the training never occurs. I changed ReplayBuffer to be a class again and it solves CartPole in <200 episodes or <10000 steps.

Other than that, I used epsilon decay to make it solve CartPole a bit faster.

[BradLarson]

@seungjaeryanlee - Oh, right, I accidentally left the class->struct conversion in there from my testing. I had been trying to see if ReplayBuffer as a class was interacting badly with the autodiff system, thus the attempt at converting it to a struct. I totally forgot to convert that back, and that makes complete sense as to why that wasn't being mutated properly.

There might be an argument in the future for converting that to a struct and then using appropriate inout semantics throughout to utilize some advantages of value types while preserving appropriate mutation, but we're good for now in having this be a class. Good catch, and glad that was that simple to fix.

Do you feel ready to convert this from a draft to a full PR for final review and submission?

[seungjaeryanlee]

Sounds good!

[dan-zheng]

Nice work, I think this PR is ready to land! I left some minor comments that can be addressed later.

[seungjaeryanlee]

I have added documentation comments, focusing mainly on the ReplayBuffer and Agent classes and the list of hyperparameters. I would love to hear your feedback on them!

[BradLarson]

Thanks for working hard on getting this ready, it's exciting to have an implementation of this in the models.

To follow on, I've created an issue in swift-apis to track lifting up the dimensionGathering() function (or whatever we end up naming it) into swift-apis and an issue in swift-models to add a little extra documentation to the Gym Readme describing the Deep Q-Network model in the same way as the other Gym models as well as indicating the additional dependencies of Matplotlib and Numpy.