declare-lab/tango2

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import numpy as np16import torch17import tqdm18 19from ...models.unet_1d import UNet1DModel20from ...pipelines import DiffusionPipeline21from ...utils import randn_tensor22from ...utils.dummy_pt_objects import DDPMScheduler23 24 25class ValueGuidedRLPipeline(DiffusionPipeline):26    r"""27    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the28    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)29    Pipeline for sampling actions from a diffusion model trained to predict sequences of states.30 31    Original implementation inspired by this repository: https://github.com/jannerm/diffuser.32 33    Parameters:34        value_function ([`UNet1DModel`]): A specialized UNet for fine-tuning trajectories base on reward.35        unet ([`UNet1DModel`]): U-Net architecture to denoise the encoded trajectories.36        scheduler ([`SchedulerMixin`]):37            A scheduler to be used in combination with `unet` to denoise the encoded trajectories. Default for this38            application is [`DDPMScheduler`].39        env: An environment following the OpenAI gym API to act in. For now only Hopper has pretrained models.40    """41 42    def __init__(43        self,44        value_function: UNet1DModel,45        unet: UNet1DModel,46        scheduler: DDPMScheduler,47        env,48    ):49        super().__init__()50        self.value_function = value_function51        self.unet = unet52        self.scheduler = scheduler53        self.env = env54        self.data = env.get_dataset()55        self.means = {}56        for key in self.data.keys():57            try:58                self.means[key] = self.data[key].mean()59            except:  # noqa: E72260                pass61        self.stds = {}62        for key in self.data.keys():63            try:64                self.stds[key] = self.data[key].std()65            except:  # noqa: E72266                pass67        self.state_dim = env.observation_space.shape[0]68        self.action_dim = env.action_space.shape[0]69 70    def normalize(self, x_in, key):71        return (x_in - self.means[key]) / self.stds[key]72 73    def de_normalize(self, x_in, key):74        return x_in * self.stds[key] + self.means[key]75 76    def to_torch(self, x_in):77        if type(x_in) is dict:78            return {k: self.to_torch(v) for k, v in x_in.items()}79        elif torch.is_tensor(x_in):80            return x_in.to(self.unet.device)81        return torch.tensor(x_in, device=self.unet.device)82 83    def reset_x0(self, x_in, cond, act_dim):84        for key, val in cond.items():85            x_in[:, key, act_dim:] = val.clone()86        return x_in87 88    def run_diffusion(self, x, conditions, n_guide_steps, scale):89        batch_size = x.shape[0]90        y = None91        for i in tqdm.tqdm(self.scheduler.timesteps):92            # create batch of timesteps to pass into model93            timesteps = torch.full((batch_size,), i, device=self.unet.device, dtype=torch.long)94            for _ in range(n_guide_steps):95                with torch.enable_grad():96                    x.requires_grad_()97 98                    # permute to match dimension for pre-trained models99                    y = self.value_function(x.permute(0, 2, 1), timesteps).sample100                    grad = torch.autograd.grad([y.sum()], [x])[0]101 102                    posterior_variance = self.scheduler._get_variance(i)103                    model_std = torch.exp(0.5 * posterior_variance)104                    grad = model_std * grad105 106                grad[timesteps < 2] = 0107                x = x.detach()108                x = x + scale * grad109                x = self.reset_x0(x, conditions, self.action_dim)110 111            prev_x = self.unet(x.permute(0, 2, 1), timesteps).sample.permute(0, 2, 1)112 113            # TODO: verify deprecation of this kwarg114            x = self.scheduler.step(prev_x, i, x, predict_epsilon=False)["prev_sample"]115 116            # apply conditions to the trajectory (set the initial state)117            x = self.reset_x0(x, conditions, self.action_dim)118            x = self.to_torch(x)119        return x, y120 121    def __call__(self, obs, batch_size=64, planning_horizon=32, n_guide_steps=2, scale=0.1):122        # normalize the observations and create  batch dimension123        obs = self.normalize(obs, "observations")124        obs = obs[None].repeat(batch_size, axis=0)125 126        conditions = {0: self.to_torch(obs)}127        shape = (batch_size, planning_horizon, self.state_dim + self.action_dim)128 129        # generate initial noise and apply our conditions (to make the trajectories start at current state)130        x1 = randn_tensor(shape, device=self.unet.device)131        x = self.reset_x0(x1, conditions, self.action_dim)132        x = self.to_torch(x)133 134        # run the diffusion process135        x, y = self.run_diffusion(x, conditions, n_guide_steps, scale)136 137        # sort output trajectories by value138        sorted_idx = y.argsort(0, descending=True).squeeze()139        sorted_values = x[sorted_idx]140        actions = sorted_values[:, :, : self.action_dim]141        actions = actions.detach().cpu().numpy()142        denorm_actions = self.de_normalize(actions, key="actions")143 144        # select the action with the highest value145        if y is not None:146            selected_index = 0147        else:148            # if we didn't run value guiding, select a random action149            selected_index = np.random.randint(0, batch_size)150 151        denorm_actions = denorm_actions[selected_index, 0]152        return denorm_actions153