[WIP] Sampled Muzero

[JosephDenman]

This is a work-in-progress implementation of Sampled Muzero that I've been working on. I figured I'd store the implementation here in case anyone else is interested in developing it. The agent learns slowly (if at all) and performs significantly worse than the vanilla MuZero agent. As I'm relatively new to these libraries, I'm out of ideas for how to debug it. One interesting discrepancy I noticed between the regular agent and the sampled agent is that the policy loss for the sampled agent initially spikes, then returns to zero, then proceeds in a logarithmic curve, whereas the regular agent's policy loss has no such initial spike. But, I don't know how to interpret this difference. It also seems that, since the policy loss does converge, that automatic differentiation is configured correctly. In that case, the question would be why the policy does not improve more than it does, which suggests that something in the tree search is misconfigured.

I'm very much interested in feedback! Thanks.

[JosephDenman]

Got it working. My problem was that the action passed to step in the environment was a tensor [i], rather than an integer i. I just had to get the integer from the tensor. Strange that the environment doesn't throw an error if the actions are incorrectly formatted.