Add UNet 1d for RL model for planning + colab

[natolambert]

Creating the PR to add back in the RL code when ready. This contains the model to run inference for Janner et. al's Diffuser, with the accompanying colab.

work in progress

TODO:

• [x] make sure old tests pass
• [x] make sure colab still works

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[ezhang7423]

Hi, is there any updates on this?

[natolambert]

Hi @ezhang7423, we have mostly paused merging this until Stable Diffusion calms down (which can be changed if there is interest!).

That being said, this PR is very stable. What are you interested in? Have you looked at the colab? Happy to discuss anything here!

[ezhang7423]

Thank you so much for the helpful response! I'm curious about whether the 1D convolution class is actually necessary- theoretically, I think you could do everything with a 2D convolution with a non square kernel. Would it be easier to just create a wrapper class on top of the existing models that does this?

I'm also wondering if you have the code for training one of these models from scratch.

[natolambert]

Ah, so this model is mostly a direct port from the original code. We would definitely welcome community contribution to help re-use the existing model components.

Part of the reason we did not is because training is complex. It involves a lot more RL infrastructure we don't really intend to support in diffusers. I think the first step would be to re-train in another repo with diffusers as a dependency.

I'm happy to discuss further if this is something you want to work on, but it's not really on my runway (because I know it will be very challenging, given my experience in RL).

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[natolambert]

Fighting the stale bot! We haven't forgotten about this, and actually moved it up the priority list today. Soon!

[natolambert]

Okay, I just updated the Colab too so it works.

Now a couple final thoughts:

1. I will add some documentation placeholders.
2. When released I will add some issues people can address to unify some similar code snippets (like Upsample1D, RearrangeDim, and Conv1D
3. I will update the colab in huggingface/notebooks when we release it and update links in the repo to that more permanent link.

[natolambert]

I agree. Was expecting some time on improving code quality, I just wanted to know where to start @patrickvonplaten. I will take a deeper look and ask questions when I'm stuck.

Most of the code is still in the form of the original.

[natolambert]

@patrickvonplaten regarding your last comment above too, yes I think this can become unet_1d. The code specific to RL was only used during training and has all been removed.

The tricky thing is that the model here uses three dimensional tensors [batch x time_dimension x state_dimension], so that may be different than some 1d use-cases. Maybe we have some optional batching? Or, the RL model becomes a wrapper of a 1d unet?

[patrickvonplaten]

If it's a very special model, then indeed let's put it in its own unet class :-) Maybe easier to begin with anyways

[bglick13]

Excited to see this is WIP! I was just going down the rabbit hole if implementing some of this myself.

The code specific to RL was only used during training and has all been removed.

Are there plans to add this back in? It would be cool to be able to replicate the ValueFunction training in the Diffusers library.

[natolambert]

@bglick13 I just don't have the time to lead it. If there is a community member that wants to lead this I am happy to advise!

[natolambert]

Making a lot of progress here, just FYI the colab will not work right now until I update the pre-trained models on the hub! I will do this in 1-2 days and post again when it's good to go!

[patrickvonplaten]

Let me know if I can help you somehow @natolambert !

[natolambert]

Thanks @patrickvonplaten, I think I am making good progress. I have a challenging step left of unifying the layer creation to be like unet_2d.py, but have been making solid progress.

[neverix]

Nice, this will be useful for porting audio and pose models

[natolambert]

@patrickvonplaten this is getting much closer. I just need to add the new 1d resnet blocks to the add_block function and it'll be very close in the API.

There is something specific to the RL model that we probably want to figure out how to remove (the model permutes the inputs to get the batching over a different dimension, which won't be in every application, so I will move it to the colab from model.).

I am going to upload a new version of the pretrained model so maybe folks can help with smoothing out features now that the big parts seem in place.

[natolambert]

@patrickvonplaten matched the get_block API and the model fusing/ddpm-unet-rl-hopper-hor128 should work again if you want to take a look.

Will do more code cleaning tomorrow (this example model is a little weird, so it will help use push the API to be flexible). I also need to change the config for the pretrained models to match the config that will be more like the UNet2D.

[natolambert]

@patrickvonplaten Any comments on adding a get_mid_block and get_output_block function? This is where our two models differ substantially.

[natolambert]

Let's make sure this plays nicely with #803!

[bglick13]

I think my implementation of value guided planning here is ready for a look. Tested on hopper and it looks good.

[natolambert]

@patrickvonplaten I merged in the changes from @bglick13 so that the work converting updates of UNet1d is more centralized and we don't both have to keep re-basing.

We decided to make it so everything fits in UNet1D and community pipelines.

Notes:

• the modifications to pipeline.py are just so Ben could debug locally / update colabs (which are updated for this change). We'll remove those before merging.
• I need to update this branch / model values to the unet midblock structure, that's why the tests are failing.

[natolambert]

@bglick13 thoughts on changing the script names to be a little clearer:

1. diffusion_generate_trajectories
2. diffusion_run_control

Or something like this?

[natolambert]

@patrickvonplaten this is really close now. I needed to do some substantial manual merging, so we need to go back and change the defaults to get your test to pass again, but that should be soon (maybe I will update this comment before you get up again).

[natolambert]

Hi @patil-suraj, sounds good happy to summarize.

This large PR is extending the functionality of UNet1D and Diffusers broadly to another applications.

The main changes are

1. Extending UNet1D to have a bit more functionality,
2. Add related 1d_blocks, which are mostly module up/downsample and some resnets that can be modified
3. Add scripts for re-producing the paper,
4. A colab for mirroring the scripts
5. Minor changes to ddpm & embeddings to re-use existing code in this model (like adding another activation function)

To address some of your specific questions:

1. I think we should re-name the existing up/downsample 1d because they aren't modules so the parameters aren't learned. They're more of functions and are different from their UNet2D counterparts
2. examples/community/pipeline.py is a placeholder file and will not be merged
3. as I said in line, diffuser is the name of the rl paper. Let's just change all mentions of diffuser -> rl so it's easier for users.

I wanted to ask if @patrickvonplaten agrees with moving to src/pipelines (I am happy to help maintain it, it's just a longer-term bet than some of the others).

And cc @bglick13.

After this, I was planning on another PR to clean up unet_1d_blocks, but I can do some of that now too.

[HuggingFaceDocBuilder]

The documentation is not available anymore as the PR was closed or merged.

[ezhang7423]

Congrats!! Thank you for all the great work.

[patrickvonplaten]

Great work - thanks a lot :-)