Not meant as an issue, but a little perplexed by feature_dim...

I noticed feature_dim is set to 50, which is quite a bottleneck from the encoding dim of 32 * 35 * 35 and the downstream hidden_dim of 1024. Very interesting. Do you think the bottleneck helps create some kind of better compression for learning?

[denisyarats]

Initially I used 256 hidden size after the encoder’s 50 dimensional output, which also works fine. Over time, however, I noticed that wider value/policy heads are useful for increasing stability of training. I suspect this is because with a wider network you have a better chance to end up in a good init position in the parameter space, so the seed won’t collapse.

[slerman12]

Yeah, that makes sense. I hadn’t thought of it that way.