This repository was archived by the owner on May 28, 2024. It is now read-only.
Repository navigation
This repository was archived by the owner on May 28, 2024. It is now read-only.
PyTorch Port #28
Description
Thread to discuss PyTorch port of PlaNet.
Firstly, I've already fixed the RNN bug here. Some things I haven't checked thoroughly are the sizes of layers/activations within the neural networks and the max transition memory size (currently 500k). I think there is probably space to parallelise some ops, will need to look into it (overshooting is the main overhead, going from O(L) with the standard KL objective to O(L^2)). I've also just switched all ReLUs to ELUs, but the results I'll post in this comment are with ReLUs (note these are incomplete because I killed runs early).
Finger Spin (can't remember if no overshooting or not but shouldn't affect results much anyway)
Walker Walk (no overshooting)
Walker Walk
Reacted by ThanksFaker
Activity
danijar commentedon Apr 23, 2019
Thanks for sharing! Comparing your two walker plots, it seems like they are quite similar in terms of performance up to the 350 episodes that the first experiment ran for. Are these training curves or testing curves (i.e. with or without action noise)?
Kaixhin commentedon Apr 23, 2019
These are training curves for a single agent/seed (I'm currently not running any test episodes - raised an issue for that: Kaixhin/PlaNet#1). In a way roughly mirrored the plots you had so my instinct was that things were probably working. Only have 1 GPU with > 4GB RAM for this and other work so unfortunately can't test much.
FYI currently moving for loops from the main code into the model. Not nearly as clean, but the only way to try and optimise things if at all.
danijar commentedon Apr 24, 2019
Maybe it gets to similar performance when training on shorter sequences but larger batch sizes. This could speed up the code quite a bit. With the fixed RSSM that you've already implemented, PlaNet now also works without overshooting and without global prior. While it would be nice to keep overshooting as it still helps in some cases, it might be a lot faster to disable it by default if it's still working well. By the way, are you encoder and decoding all images of the current chunk at once or separately for each time step? This can make a big difference.
Kaixhin commentedon Apr 24, 2019
Yes - shorter sequences and no overshooting would definitely speed things up in general. I also suspect that reducing the overshooting length could be reduced even if enabled - have you done any experiments on this (e.g., at the chunk length it's far longer than the planning length)? Over a year ago I tried to use a Deep Kalman Filter to do model-based RL with long time dependencies and the prior's predictions just collapsed to the mode, so overshooting sounded like a potential solution to this problem I had.
My current code does encoding and decoding separately per timestep, but the version I'm working on now will do it for a chunk at once.
piojanu commentedon Apr 26, 2019
@Kaixhin amazing work! I have a question, what does it mean: "the prior's predictions just collapsed to the mode"? Specifically, what do you understand by "mode"? Could you post some references where I could read more on the topic?
Kaixhin commentedon Apr 26, 2019
@piojanu thank you :) It's a bit off-topic but briefly I did some work on MBRL with latent space planning called DELIP, and in some experiments that I didn't get to work reliably I was trying to predict trajectories of about 100 timesteps in the future. In this case the trajectories happened to fall into a couple of different clusters, and the learned model predicted the most common ones but completely missed ones in the smaller clusters. Well, I did train some models to get it right, but I wanted to find a more reliable setup before publishing results. Unfortunately missed the window to do so but hoping I could pick up something similar again using PlaNet as a strong starting point.
Kaixhin commentedon Apr 27, 2019
Update on my refactor - tqdm is predicting about 1 week (I said the previous version was about 2 weeks, but that was an upper bound). Had a look at GPU comparisons - this one is giving a small multiplier on using a V100 vs. a 1080Ti, so with a dedicated headless server and a better GPU this timing should be a bit less too. Changes include:
MUJOCO_GL=egl(not sure this makes much difference for me though)Started a run with Walker Walk with default hyperparams, fingers crossed it replicates what you have.
piojanu commentedon Apr 27, 2019
@Kaixhin thanks for explanation, now I get it 👌 The paper though looks great. I'm especially interested in how you adopted the Monte Carlo planner to learned model, because I'll do something similar in PlaNet in next steps! I'll read it for sure, thanks 😄
Kaixhin commentedon Apr 28, 2019
@danijar just checking a few details on the losses. In general I try to sum over elements (e.g. MSE loss on pixels, KL loss on 30D Gaussian), average over batch, sum over time, but perhaps this is a bit off.
danijar commentedon Apr 29, 2019
Exactly, losses are summed across the feature dimension and averaged across batch and time. Your summary of free nats sounds correct to me.
Kaixhin commentedon Apr 29, 2019
Going back to the loss in the paper (looked at #19, where it seems like the implementation details follow the paper but not explicitly so):
All losses (observation NLL + reward NLL + KL) should actually be summed over time? On top of this the KL losses are normalised by the maximum chunk length (so 1/D factor) - not sure I can pick up from the paper whether the global KL should also be normalised by this as well? Though with β1 = D this normalisation is effectively undone for the standard KL loss (β>1 = 1 means no change for the overshooting part).
danijar commentedon May 1, 2019
Yes, all losses are averaged over both batch and time dimensions, including the global divergence.
Kaixhin commentedon May 1, 2019
Sorry for not being clear - I've tried to write up the loss informally to check - let me know if this matches what you have:
14 remaining items
Kaixhin commentedon May 21, 2019
I've got some results and pretrained models in a release, with the rest to come. Everything looks good, except
ball_in_cup-catchis pretty high variance. Something I forgot was to add compression to 5 bits (I used 8 bits) - so I can try this. Any idea how much difference this might make?Edit: Also, is there some sort of exploration schedule?
piojanu commentedon May 23, 2019
@Kaixhin there is code for it e.g. here it's applied:
planet/planet/control/mpc_agent.py
Line 75 in 0c6f7d3
It's not used as you can see (
exploration_rampsis 0):planet/planet/tools/schedule.py
Lines 33 to 36 in 0c6f7d3
danijar commentedon May 24, 2019
@Kaixhin The plots looks great! I've also heard several people using your implementation of PlaNet already, so it seem to have a good impact. As @piojanu says (you have a really good idea of the code now haha!) I'm not using any schedules, neither for learning rate nor for exploration. The features are in the code though and I've played around with them at some point. Do you want to add the figures on the release page for completeness?
Kaixhin commentedon May 25, 2019
Yep great to hear that people have been using my port 😄 I'll add the figures on the release page once I've done a final check on whether those are the settings that should be used (see below).
Still looking at
ball_in_cup-catch. Added 5-bit quantisation (including random uniform noise for the appropriate bin size; also centred the data by adding -0.5 just to finish making preprocessing the same). Unfortunately results look even worse (below without and with the learning rate schedule you originally had). There's a chance that I've messed something up with experience replay storage (because I'm quantising back to bytes for smaller storage), but I'm just going to revert back to the simple 8-bit version without noise and rerun to see if I get something like the previous results, which were at least roughly in the right area even if the variance on that env is very high.danijar commentedon May 25, 2019
I see. I'm plotting the median rather than the mean in the PlaNet paper, because that's what the dm_control paper does. In general, the median makes sense to me but it can be problematic for tasks with basically binary reward like ball in cup. However, the max values in your plots are definitely too low. Are you giving the agent 3 free nats and is this without overshooting?
If you want to try the preprocessing, it should be easy to check whether it's correct by applying the inverse transform and looking at the resulting images. If they look good, and the preprocessed values are in range -0.5 to 0.5, then it's likely to be correct.
Kaixhin commentedon May 27, 2019
Yes these results are from using the camera ready hyperparameters. I think there was an issue in my discretisation for saving observations in experience replay memory (I had
preprocessfine, but notpostprocess). Might be OK now (seems to learn quickly, but the policy performance is still a bit unstable)?danijar commentedon May 27, 2019
It might be fine, although it's still a bit surprising. There might be a difference in how we read the data, or neural network initialization etc because of difference between PyTorch and TensorFlow. If you can run multiple seeds, I'd try plotting the median aggregating over seeds and multiple episodes. That way, we'll have comparable figures. But since the performance matches quite closely on the other tasks I wouldn't worry too much about it.
Kaixhin commentedon May 30, 2019
I've considered network initialisation, but for this I doubt it would be a big factor if you're not doing something custom. PyTorch initialises both weights and biases with something similar to Kaiming init, and I believe TensorFlow uses truncated normal for weights and zeroed biases; I'm not sure it's worth implementing truncated normal for PyTorch.
I've added median plotting (dark blue line) and tried a different seed and it looks alright. Median and mean seem to track each other pretty well for this. Getting more experiments on the other domains to make sure that performance is similar, so fingers crossed.
danijar commentedon Jun 22, 2019
Sorry for the late reply. How many episodes/seeds is the median taken over? It's surprising that the median is somewhat close to 500 because I'd expect the agent to either get the ball in (close to 1000) or not (close to 0). So the median should be close to one of these values most of the time, if aggregating over enough episodes/seeds. On the other hand, as I said earlier, I wouldn't worry too much about this, since the performance looks very nice on the other tasks.
Kaixhin commentedon Jun 22, 2019
This is just 1 seed (don't have resources to run any more experiments on this sadly). Stats are over 10 evaluation episodes, and I've seen many where some miss once or twice but do get it in eventually. In the releases I've noted performance replication.
As you said, this hopefully isn't indicative of any problems given that all the other results look fine, so I think I'll close this issue up here. Thanks for all the help!
danijar commentedon Jun 23, 2019
Sounds good, thanks for the great work!
mctigger commentedon Jul 9, 2020
Hey @Kaixhin
I am wondering if you were able to reproduce the results from the paper with your implementation? If I understand the graphs correctly, the results in the paper indicate much faster convergence than the PyTorch port?
Kaixhin commentedon Jul 9, 2020
@mctigger from release 1.2,
cartpole-swingup,reacher-easy,finger-spinandwalker-walkare about the same, thoughball_in_cup-catchis a bit unstable andcheetah-runis slow. The latest version of the code now includes action clipping (Kaixhin/PlaNet#22), which wasn't specified in the paper but is used in this codebase, so with that it should ideally be comparable (but I don't have resources to test atm).mctigger commentedon Jul 9, 2020
Thank you for the quick response!