| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import torch |
| from .batch import TensorBatch as DataProto |
|
|
| def msp_reward_fn(self, batch: DataProto, pixels, return_reward_tensor=True, save_pred=False, pixels_before_repeat=None): |
| if pixels_before_repeat is None: |
| pixels_before_repeat = pixels |
|
|
| batch_size = batch.batch["responses"].shape[0] |
| segment_length = self.config.data.video.segment_length |
| tokens_per_frame = self.config.processor.tokens_per_frame |
| action_dim = self.config.processor.action_dim |
|
|
| output_tokens = batch.batch["responses"].reshape(batch_size, segment_length - 1, tokens_per_frame + action_dim) |
| output_tokens = output_tokens[:, :, :tokens_per_frame] |
| output_tokens = output_tokens.clamp(0, self.config.processor.visual_token_num - 1).long() |
|
|
| ctx_tokens = batch.batch["ctx_tokens"] |
| output_tokens = output_tokens.reshape(batch_size, segment_length - 1, tokens_per_frame) |
|
|
| detokenize_output = self.tokenizer_wg.detokenize( |
| DataProto.from_single_dict({"tokens": output_tokens, "ctx_tokens": ctx_tokens}), |
| |
| DataProto.from_single_dict({"dummy": torch.zeros((batch_size, 1))}, meta_info={'lpips': True, 'recon': self.config.trainer.reward_fn}), |
| ) |
|
|
| if 'recon_loss' in detokenize_output.batch.keys(): |
| recon_loss = detokenize_output.batch['recon_loss'] |
| else: |
| pred = detokenize_output.batch['pixels'].clamp(0.0, 1.0)[:, 1:] |
| real = pixels[:, 2:] |
| if self.config.trainer.reward_fn == 'mse': |
| recon_loss = torch.mean((real - pred)**2, dim=(2, 3, 4)) |
| elif self.config.trainer.reward_fn == 'mae': |
| recon_loss = torch.mean(torch.abs(real - pred), dim=(2, 3, 4)) |
| else: |
| raise NotImplementedError(f"Unsupported reward function: {self.config.trainer.reward_fn}") |
|
|
| if 'perceptual_loss' in detokenize_output.batch.keys(): |
| perceptual_loss = detokenize_output.batch['perceptual_loss'] |
| else: |
| perceptual_loss = self.tokenizer_wg.perceptual_loss(DataProto.from_single_dict({"real": real, "pred": pred})) |
| perceptual_loss = perceptual_loss.batch['perceptual_loss'] |
|
|
| if self.config.trainer.msp_reward_aggregate == "mean": |
| loss = (recon_loss + perceptual_loss).mean(-1) |
| elif self.config.trainer.msp_reward_aggregate == "last": |
| loss = (recon_loss + perceptual_loss)[:, -1] |
| elif self.config.trainer.msp_reward_aggregate == "discount": |
| discount = 0.9 |
| weight = discount ** torch.arange(recon_loss.shape[1], device=recon_loss.device) |
| loss = (recon_loss + perceptual_loss) * weight.unsqueeze(0) |
| loss = loss.sum(-1) / weight.sum() |
| if not return_reward_tensor: |
| return loss |
|
|
| print("loss", loss.mean().item(), flush=True) |
| print("recon_loss", recon_loss.mean().item(), flush=True) |
| print("perceptual_loss", perceptual_loss.mean().item(), flush=True) |
| if save_pred: |
| pred = detokenize_output.batch['pixels'].clamp(0.0, 1.0)[:, 1:] |
| real = pixels[:, 2:] |
| for i in range(1): |
| |
| real_traj = torch.cat([real[i:i+1, j] for j in range(real.shape[1])], dim=-1) |
| pred_traj = torch.cat([pred[i:i+1, j] for j in range(pred.shape[1])], dim=-1) |
| plot_img(torch.cat([ |
| real_traj.float(), |
| pred_traj.float(), |
| torch.abs(real_traj-pred_traj).float(), |
| ], dim=-2), f"{self.config.trainer.experiment_name}-{i}") |
|
|
| reward_tensor = torch.zeros_like(batch.batch['responses'], dtype=torch.float32) |
| for i in range(len(batch)): |
| data_item = batch[i] |
| prompt_ids = data_item.batch['prompts'] |
| prompt_length = prompt_ids.shape[-1] |
| valid_response_length = data_item.batch['attention_mask'][prompt_length:].sum().long().item() |
| reward_tensor[i, valid_response_length - 1] = -loss[i].item() |
|
|
| return reward_tensor, { |
| "critic/recon_loss/mean": recon_loss.mean().item(), |
| "critic/perceptual_loss/mean": perceptual_loss.mean().item(), |
| } |
|
|
