斯坦福新成果：用「4D时空坐标系」重构机器人视觉认知，攻克跨视角一致性难题！

class Decoder(nn.Module):    def forward(self, z, **kwargs):            # assert z.shape[1:] == self.z_shape[1:]            self.last_z_shape = z.shape
            # timestep embedding            temb = None
            # z to block_in            h = self.conv_in(z)
            # middle            h = self.mid.block_1(h, temb, **kwargs)            h = self.mid.attn_1(h, **kwargs)            h = self.mid.block_2(h, temb, **kwargs)
            # upsampling            for i_level in reversed(range(self.num_resolutions)):                for i_block in range(self.num_res_blocks + 1):                    h = self.up[i_level].block[i_block](h, temb, **kwargs)                    if len(self.up[i_level].attn) > 0:                        h = self.up[i_level].attn[i_block](h, **kwargs)                if i_level != 0:                    h = self.up[i_level].upsample(h)
            # end            if self.give_pre_end:                return h
            h = self.norm_out(h)            h = nonlinearity(h)            h = self.conv_out(h, **kwargs)            if self.tanh_out:                h = torch.tanh(h)            return h

class DiffusionEngine(pl.LightningModule):    def get_mv_input(self, batch):            # get images from different views            return torch.cat([batch[input_key] for input_key in self.input_mv_keys], dim=0)
    def encode_first_stage(self, x):            if len(x.shape) == 5:                x = rearrange(x, "b t c h w -> (b t) c h w")
            n_samples = default(self.en_and_decode_n_samples_a_time, x.shape[0])            n_rounds = math.ceil(x.shape[0] / n_samples)            all_out = []            with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):                for n in range(n_rounds):                    out_1 = self.first_stage_pointmap_model.encode(                        x[n * n_samples: (n + 1) * n_samples, :3]                    )                    out_2 = self.first_stage_color_model.encode(                        x[n * n_samples: (n + 1) * n_samples, 3:]                    )                    out = torch.cat([out_1, out_2], dim=1)                    all_out.append(out)            z = torch.cat(all_out, dim=0)            z = self.scale_factor * z            return z
    def forward(self, x, batch):            with torch.autocast("cuda", enabled=not self.disable_loss_fn_autocast):                loss = self.loss_fn(self.model, self.denoiser, self.conditioner, x, batch)
            loss_mean = loss.mean()            loss_dict = {"loss": loss_mean}            return loss_mean, loss_dict
    def shared_step(self, batch: Dict) -> Any:            x = self.get_mv_input(batch)            x = self.encode_first_stage(x)            batch["global_step"] = self.global_step            loss, loss_dict = self(x, batch)            return loss, loss_dict
    def training_step(self, batch, batch_idx):            batch['num_video_frames'] = batch['num_video_frames'][0]            loss, loss_dict = self.shared_step(batch)            loss_dict['loss'] = loss            ...            return loss_dict