双系统的训练

[markAKB]

在训练时使用了traj_images提取到的特征，进行噪声的预测。而加噪轨迹基于traj_poses，两者在时间上是一一对应的。请问我的理解有问题吗？如果是这样，是不是会训练中会出现未来帧的信息泄露的问题？

以下是我阅读到的训练时forward和dataset加载的代码：

InternNav/internnav/model/basemodel/internvla_n1/internvla_n1.py

Line 255 in 0f82252

latents = torch.cat([memory_tokens, traj_hidden_states], dim=1)

InternNav/internnav/dataset/internvla_n1_lerobot_dataset.py

Line 1127 in 0f82252

traj_poses_gt.append(torch.tensor(rel_pose_resample))

[kellyiss]

训练时是没有未来帧出现的 只有历史和当前帧 而预测的trajectory是未来的轨迹

[markAKB]

训练时是没有未来帧出现的 只有历史和当前帧 而预测的trajectory是未来的轨迹

InternNav/internnav/model/basemodel/internvla_n1/internvla_n1.py

Line 255 in 0f82252

latents = torch.cat([memory_tokens, traj_hidden_states], dim=1)

        if 'nextdit' in self.get_system1_type():
                if 'async' in self.get_system1_type():
                    cur_images = traj_images.flatten(0, 1)
                    pix_goal_images = traj_images[:, 0:1].repeat(1, traj_images.size(1), 1, 1, 1).flatten(0, 1)
                    bsz = cur_images.size(0)
                    images_dp = torch.stack([pix_goal_images, cur_images], dim=1).permute(0, 1, 4, 2, 3)
                    images_dp_norm = (images_dp - self._resnet_mean) / self._resnet_std

                    images_dp_feat = (
                        self.get_model()
                        .rgb_model.get_intermediate_layers(images_dp_norm.flatten(0, 1))[0]
                        .unflatten(dim=0, sizes=(bsz, -1))
                    )

                    memory_feat = self.get_model().memory_encoder(
                        images_dp_feat.flatten(1, 2)
                    )  # [bs*select_size,512,384]
                    memory_feat = torch.cat([images_dp_feat.flatten(1, 2), memory_feat], dim=-1)
                    memory_tokens = self.get_model().rgb_resampler(memory_feat)

                    traj_hidden_states = self.get_model().cond_projector(traj_hidden_states)
                    latents = torch.cat([memory_tokens, traj_hidden_states], dim=1)
                else:
                    traj_hidden_states = self.get_model().cond_projector(traj_hidden_states)
                    latents = traj_hidden_states

这里的memory_tokens是根据traj_images提取得到的，而traj_images在dataset里对应的是未来帧？

        if self.pixel_goal_only:
            goal_len = end_frame_id - start_frame_id - 1
            interval = 2
            frame_ids = np.arange(0, goal_len, interval)
            max_len = 12
            traj_images = torch.tensor(np.stack([np.asarray(timg.resize((224, 224))) for timg in traj_images])) / 255.0
            if len(frame_ids) > max_len:
                interval = int(np.ceil(goal_len / max_len))
                frame_ids = np.arange(0, goal_len, interval)

            traj_poses_gt = []
            for cid in frame_ids:
                discrete_traj_pose = get_trajectory_relative_to_frame(pose[cid:], camera_deg=pitch_2)
                rel_trajectory, rel_pose_resample = interpolate_and_resample_trajectory(
                    discrete_traj_pose, self.predict_step_num
                )
                rel_pose_resample = clip_or_pad(rel_pose_resample, self.predict_step_num)
                traj_poses_gt.append(torch.tensor(rel_pose_resample))

            data_dict["traj_images"] = traj_images[:goal_len][::interval]
            data_dict["traj_depths"] = torch.stack(traj_depths[:goal_len][::interval])
            data_dict["traj_poses"] = torch.stack(traj_poses_gt)

[kellyiss]

在异步的训练中 用到了pixel goal image （x_0） 和 current image (x_t). 我们每次是将 traj_images 中的x_t取出来 去预测 x_t+1 之后的轨迹。 这里是一个batch的操作 memory token只有两张图的feat

[markAKB]

在异步的训练中 用到了pixel goal image （x_0） 和 current image (x_t). 我们每次是将 traj_images 中的x_t取出来 去预测 x_t+1 之后的轨迹。 这里是一个batch的操作 memory token只有两张图的feat

提到的逻辑实现是在如下代码处吗？
我打印了traj_images的shape为[4,8,224,224,3]，bs=4，包含了未来8帧图片。
而取出current image (x_t)是通过cur_images = traj_images.flatten(0, 1)实现的，cur_images 的shpe为[32,224,224,3]，并不是如提到的取出x_t单独的这一帧？

        if 'nextdit' in self.get_system1_type():
                if 'async' in self.get_system1_type():
                    cur_images = traj_images.flatten(0, 1)
                    pix_goal_images = traj_images[:, 0:1].repeat(1, traj_images.size(1), 1, 1, 1).flatten(0, 1)
                    bsz = cur_images.size(0)
                    images_dp = torch.stack([pix_goal_images, cur_images], dim=1).permute(0, 1, 4, 2, 3)
                    images_dp_norm = (images_dp - self._resnet_mean) / self._resnet_std

                    images_dp_feat = (
                        self.get_model()
                        .rgb_model.get_intermediate_layers(images_dp_norm.flatten(0, 1))[0]
                        .unflatten(dim=0, sizes=(bsz, -1))
                    )

                    memory_feat = self.get_model().memory_encoder(
                        images_dp_feat.flatten(1, 2)
                    )  # [bs*select_size,512,384]
                    memory_feat = torch.cat([images_dp_feat.flatten(1, 2), memory_feat], dim=-1)
                    memory_tokens = self.get_model().rgb_resampler(memory_feat)

                    traj_hidden_states = self.get_model().cond_projector(traj_hidden_states)
                    latents = torch.cat([memory_tokens, traj_hidden_states], dim=1)
                else:
                    traj_hidden_states = self.get_model().cond_projector(traj_hidden_states)
                    latents = traj_hidden_states

[kellyiss]

32 是在batch那一维 这里就是future trajectory 每一步都是xt 组成一个batch 训练了

[markAKB]

32 是在batch那一维 这里就是future trajectory 每一步都是xt 组成一个batch 训练了

我设置的bs是=4，traj_images.flatten(0, 1)这个操作是把bs和future traj_len这两个维度合并了，future traj_len=8（未来的图片是8帧），所以是维度变为了32

[kellyiss]

InternNav/internnav/model/basemodel/internvla_n1/internvla_n1.py

Line 239 in 0f82252

images_dp = torch.stack([pix_goal_images, cur_images], dim=1).permute(0, 1, 4, 2, 3)

可以打印这里的size看下 设的bs是system 2的batch size 这里为了训到异步 system 1的batch size 是bs * future traj_len
对于每个sample来说 只有两张图 pixel goal image 和 cur image

[markAKB]

InternNav/internnav/model/basemodel/internvla_n1/internvla_n1.py

Line 239 in 0f82252

images_dp = torch.stack([pix_goal_images, cur_images], dim=1).permute(0, 1, 4, 2, 3)

可以打印这里的size看下 设的bs是system 2的batch size 这里为了训到异步 system 1的batch size 是bs * future traj_len
对于每个sample来说 只有两张图 pixel goal image 和 cur image

明白了，很优雅的代码，感谢耐心解答。