model/flow/mlp_shortcut.py

# MIT License

# Copyright (c) 2025 ReinFlow Authors

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


# MIT License

# Copyright (c) 2024 Kevin Frans

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# The description of ShortCutFlowMLP is translated from Kevin Fran's One Step Diffusion via Short Cut Models 
# and revised by ReinFlow Authors and Collaborators.
# The ShortCutFlowViT scipt is extended from Diffusion Policy Policy Optimization's implementation. 
# NoisyShortCutFlowMLP and NoisyShortCutFlowViT are composed by ReinFlow Authors.
  
import torch
import torch.nn as nn
import numpy as np
import logging
from copy import deepcopy
from torch import Tensor
from diffusion_policy.model.common.mlp import MLP, ResidualMLP
from diffusion_policy.model.diffusion.positional_embedding import SinusoidalPosEmb
from diffusion_policy.model.common.modules import SpatialEmb, RandomShiftsAug
from diffusion_policy.model.common.vit import VitEncoder
from diffusion_policy.model.flow.mlp_flow import NoisyFlowMLP, ExploreNoiseNet
log = logging.getLogger(__name__)
import einops
from typing import Tuple, List

class ShortCutFlowMLP(nn.Module):
    def __init__(
        self,
        horizon_steps,
        action_dim,
        cond_dim,
        td_emb_dim=16,  # Embedding dimension for time and step
        mlp_dims=[256, 256], # hidden layers of the velocity head.
        cond_mlp_dims=None, # the hidden dimensions and output dimension of condition embedder. 
        activation_type="Mish", # different from reflow (SiLU is better for dense nets)
        out_activation_type="Identity",
        use_layernorm=False,
        residual_style=False,
        embed_combination_type='add' #, multiply or concate.   cond_embed + td_embed or cond_embed x td_embed or [cond_embed, td_embed]
        ):
        super().__init__()
        self.td_emb_dim = td_emb_dim # for both time and step
        self.act_dim_total = action_dim * horizon_steps
        self.horizon_steps = horizon_steps
        self.action_dim = action_dim
        self.cond_dim=cond_dim
        self.activation_type=activation_type
        self.out_activation_type=out_activation_type
        self.time_embed_activation=nn.Mish() # nn.SiLU() maybe better but for fair comparison with reflow and diffusion we use Mish. 
        self.use_layernorm=use_layernorm
        self.residual_style=residual_style
        candidate_embed_combination_types=['add', 'multiply', 'concate']
        if embed_combination_type not in candidate_embed_combination_types:
            raise ValueError(f"embed_combination_type must be one of {candidate_embed_combination_types} but received {embed_combination_type}!")
        self.embed_combination_type=embed_combination_type
        
        # time t and step share an input embedding
        self.map_noise = SinusoidalPosEmb(td_emb_dim)
        # MLP to process concatenated t and step embeddings
        self.t_emb = nn.Sequential(
            nn.Linear(2 * td_emb_dim, td_emb_dim),
            self.time_embed_activation,
            nn.Linear(td_emb_dim, td_emb_dim)
        )
        
        # Condition embedding
        if cond_mlp_dims:
            self.cond_emb = MLP(
                [cond_dim] + cond_mlp_dims,
                activation_type=activation_type,
                out_activation_type="Identity",
            )
            self.cond_enc_dim = cond_mlp_dims[-1]
        else:
            self.cond_enc_dim = cond_dim
        if embed_combination_type in ['add', 'multiply'] and td_emb_dim !=self.cond_enc_dim: 
            raise ValueError(f"To add or multiply td_embed with cond_embed you must make td_emb_dim={td_emb_dim} == self.cond_enc_dim={self.cond_enc_dim}")
        
        
        # velocity head
        model = ResidualMLP if residual_style else MLP
        if self.embed_combination_type =='concate':
            input_dim = action_dim * horizon_steps + self.cond_enc_dim + td_emb_dim #(s,a, t-dt)
        elif self.embed_combination_type =='add' or 'multiply':
            input_dim = action_dim * horizon_steps + td_emb_dim
        else:
            raise ValueError(f"Unsupported embed_combination_type={self.embed_combination_type}")
        self.vel_head = model(
            [input_dim] + mlp_dims + [self.act_dim_total],
            activation_type=activation_type,
            out_activation_type=out_activation_type,
            use_layernorm=use_layernorm,
        )

    def forward(
        self,
        action: Tensor,
        time: Tensor,
        dt: Tensor,
        cond: dict,
        output_embedding=False
        ):
        """
        Inputs:
            action: (B, Ta, Da) - Current action trajectory
            time: (B,) - Current noise level t
            cond: (B, Do) - Condition (e.g., flattened state)
            dt: (B,) - Step size

        Outputs:
            velocity: (B, Ta, Da) - Predicted velocity
        """
        B, Ta, Da = action.shape

        # Flatten action
        action_flat = action.view(B, -1)
        
        # Embed time t and dt separately, then concatenate and feed to the same MLP to squeeze the dimension back to emb_dim
        t_emb = self.map_noise(time.view(B, 1)).view(B, self.td_emb_dim)
        dt_emb = self.map_noise(dt.view(B, 1)).view(B, self.td_emb_dim)
        td_emb = self.t_emb(torch.cat([t_emb, dt_emb], dim=1))

        # Embed condition and add to time-step embedding
        state = cond["state"].view(B, -1)
        cond_emb = self.cond_emb(state) if hasattr(self, "cond_emb") else state
            
        if self.embed_combination_type=='add':# we use add to reduce dimension
            emb = td_emb + cond_emb
        elif self.embed_combination_type=='multiply':# we use add to reduce dimension while preserving nonlinearity
            emb = td_emb * cond_emb
        elif self.embed_combination_type =='concate': # separate the influences of td_embd and cond_emb
            emb=torch.cat([td_emb, cond_emb], dim=-1)
        # Predict velocity
        vel_flat = self.vel_head(torch.cat([action_flat, emb], dim=-1))
        if output_embedding:
            return vel_flat.view(B, Ta, Da), td_emb, cond_emb
        return vel_flat.view(B, Ta, Da)
    
    def sample_action(self,cond:dict,inference_steps:int,clip_intermediate_actions:bool,act_range:List[float], z:Tensor=None,save_chains:bool=False):
        """
        simply return action via integration (Euler's method). the initial noise could be specified. 
        when `save_chains` is True, also return the denoising trajectory.
        """
        B = cond['state'].shape[0]
        device=cond['state'].device

        x_hat:Tensor=z if z is not None else torch.randn(B, self.horizon_steps, self.action_dim, device=device)
        if save_chains:
            x_chain=torch.zeros((B, inference_steps+1, self.horizon_steps, self.action_dim), device=device)
        dt = (1 / inference_steps) * torch.ones_like(x_hat, device=device)
        steps = torch.linspace(0, 1-1/inference_steps, inference_steps, device=device).repeat(B, 1)
        for i in range(inference_steps):
            t = steps[:, i]
            dt_batch = (1 / inference_steps) * torch.ones(B, device=device)
            vt = self.forward(action=x_hat, time=t, dt=dt_batch, cond=cond, output_embedding=False)
            x_hat += vt * dt
            if clip_intermediate_actions or i == inference_steps-1: # always clip the output action. appended by Tonghe on 04/25/2025
                x_hat = x_hat.clamp(*act_range)
            if save_chains:
                x_chain[:, i+1] = x_hat
        if save_chains:
            return x_hat, x_chain
        return x_hat


class ShortCutFlowViT(nn.Module):
    """With ViT backbone and Transformer-based shortcut flow
    
    
    **Architecture**:
    
    camera pixels -> aug-> backbone-> visual_feature - | 
                                                      cat->cond_embed->cond_embedding->|
    proprioception-> prop embedder -> prop_embedding - |                               |
                                                                                       + or * --> cond_td_embedding-|
            t     ->               -> t_embedding ->                                   |                            |
                        map_noise                     td_embed     -->   td_embedding->|                            |
           step   ->               -> dt_embedding->                                                           cat --> vel_head --> vel
                                                                                                                | 
          action  -->   (omitted)                                                         -->  act_embedding   -|
                        (projection + positional embedding)
    """
    def __init__(
        self,
        backbone:VitEncoder,    # VitEncoder instance
        action_dim,
        horizon_steps,
        prop_dim,               # proprioception dimension
        img_cond_steps=1,
        td_emb_dim=16,  # Embedding dimension for time and step
        # d_model=384,  # omitted
        # n_heads=6,
        # depth=12,
        mlp_dims=[256,256], # the hidden dimensions of output velocity head
        cond_mlp_dims=None, # the hidden dimensions and output dimension of condition embedder. 
        activation_type="Mish", # instead of SiLU() maybe better for deep nets.
        out_activation_type="Identity",
        use_layernorm=False,
        residual_style=False,
        dropout=0.0,
        visual_feature_dim=128, # overload spatial embed when specified.
        num_img=1,
        augment=False,
        spatial_emb=0,
        embed_combination_type='add'  # 'add', 'multiply', or 'concate'
    ):
        super().__init__()
        
        # Action chunk
        self.action_dim = action_dim
        self.horizon_steps = horizon_steps
        self.act_dim_total = action_dim * horizon_steps
        
        # Historical proprioception and visual inputs
        self.prop_dim = prop_dim
        self.img_cond_steps = img_cond_steps
        
        # How to combine timestep and condition embeddings.
        candidate_embed_combination_types = ['add', 'multiply', 'concate']
        if embed_combination_type not in candidate_embed_combination_types:
            raise ValueError(f"embed_combination_type must be one of {candidate_embed_combination_types}, got {embed_combination_type}")
        self.embed_combination_type = embed_combination_type
        
        # Transformer dimension (omitted)
        # self.d_model = d_model
        
        # Action embeddings: projection and action chunk positional embedding (omitted)
        # self.x_proj = nn.Linear(action_dim, d_model)
        # self.pos_emb = PositionalEmbedding(d_model)
        
        # Time-step embeddings
        self.td_emb_dim = td_emb_dim
        self.map_noise = SinusoidalPosEmb(td_emb_dim)   # Shared embedding for time t and step size d
        self.time_embed_activation=nn.Mish()            # nn.SiLU() maybe better but for fair comparison with reflow and diffusion we use Mish. 
        self.td_emb = nn.Sequential(
                nn.Linear(2 * td_emb_dim, td_emb_dim),
                self.time_embed_activation,
                nn.Linear(td_emb_dim, td_emb_dim)
            )
        
        # Condition embedding
        if cond_mlp_dims:# add transform to the state 
            self.prop_emb = MLP(
                [prop_dim] + cond_mlp_dims,
                activation_type=activation_type,
                out_activation_type="Identity",
            )
            self.prop_embed_dim = cond_mlp_dims[-1]
        else: # just us the state itself, without transforms. 
            self.prop_embed_dim = prop_dim
        
        # Visual backbone and augmentation
        self.backbone = backbone
        self.num_img = num_img
        self.augment = augment
        if augment:
            self.aug = RandomShiftsAug(pad=4)
        # Visual feature compression
        if spatial_emb > 0:
            assert spatial_emb > 1, "spatial_emb must be > 1"
            if num_img == 2:
                self.compress1 = SpatialEmb(
                    num_patch=self.backbone.num_patch,
                    patch_dim=self.backbone.patch_repr_dim,
                    prop_dim=prop_dim,
                    proj_dim=spatial_emb,
                    dropout=dropout,
                )
                self.compress2 = deepcopy(self.compress1)
            elif num_img == 1:
                self.compress = SpatialEmb(
                    num_patch=self.backbone.num_patch,
                    patch_dim=self.backbone.patch_repr_dim,
                    prop_dim=prop_dim,
                    proj_dim=spatial_emb,
                    dropout=dropout,
                )
            else:
                raise NotImplementedError(f"num_img={num_img} not supported (only 1 or 2)")
            self.visual_feature_dim = spatial_emb * num_img
        else:
            self.visual_feature_dim = visual_feature_dim
            self.compress = nn.Sequential(
                nn.Linear(self.backbone.repr_dim, visual_feature_dim),
                nn.LayerNorm(visual_feature_dim),
                nn.Dropout(dropout),
                nn.ReLU(),
            )
        self.visuomotor_feature_dim = self.visual_feature_dim + self.prop_embed_dim
        
        if embed_combination_type in ['add', 'multiply']: 
            # compress visuomotor information to the same size of time embedding.
            self.cond_embed=nn.Sequential(
                nn.Linear(self.visuomotor_feature_dim, td_emb_dim*2),
                nn.ReLU(),
                nn.Linear(td_emb_dim*2, td_emb_dim),
            )
            self.cond_enc_dim=td_emb_dim
        else:
            self.cond_enc_dim=self.visuomotor_feature_dim
        
        # Transformer middle blocks (omitted)
        # self.transformer_blocks = nn.ModuleList([
        #     ShortcutDiTBlock(d_model, n_heads, dropout) for _ in range(depth)
        # ])
        

        # velocity head
        vel_head_model = ResidualMLP if residual_style else MLP
        if self.embed_combination_type =='concate':
            input_dim = action_dim * horizon_steps + self.cond_enc_dim + td_emb_dim     #(s, a, t-dt)
        elif self.embed_combination_type =='add' or 'multiply':
            input_dim = action_dim * horizon_steps + self.cond_enc_dim                  
        else:
            raise ValueError(f"Unsupported embed_combination_type={self.embed_combination_type}")
        output_dim = action_dim * horizon_steps
        self.vel_head = vel_head_model(
            [input_dim] + mlp_dims + [output_dim],
            activation_type=activation_type,
            out_activation_type=out_activation_type,
            use_layernorm=use_layernorm,
        )
    
    def forward(
        self,
        action,
        time,
        d,
        cond,
        output_embedding=False,
    ):
        """
        Inputs:
            action: (B, Ta, Da) - Action trajectories
            time: (B,) or float - Flow time
            d: (B,) or float - Step size
            cond: dict with keys 'state' and 'rgb'
                state: (B, To, Do) - Proprioceptive states
                rgb: (B, To, C, H, W) - RGB images
            output_embedding: whether also return td_embedding and condition embedding
        Outputs:
            velocity: (B, Ta, Da) - Predicted velocities
        """
        B, Ta, Da = action.shape
        _, T_rgb, C, H, W = cond["rgb"].shape
        
        # flatten chunk
        action_embed = action.view(B, -1)
        # (action transform omitted)
        # Project action chunk and add positional embeddings
        # x = self.x_proj(action)  # (B, Ta, d_model)
        # pos_emb = self.pos_emb(torch.arange(Ta, device=device))
        # x = x + pos_emb[None, :]
        
        # Embed time t and dt separately, then concatenate and feed to the same MLP to squeeze the dimension back to `td_emb_dim`
        t_emb = self.map_noise(time.view(B, 1)).view(B, self.td_emb_dim)
        d_emb = self.map_noise(d.view(B, 1)).view(B, self.td_emb_dim)
        td_emb = self.td_emb(torch.cat([t_emb, d_emb], dim=1))
        
        # Embed proprioceptive states
        state = cond["state"].view(B, -1)
        prop_emb = self.prop_emb(state) if hasattr(self, "prop_emb") else state
        
        # Process visual inputs (augmentation + compression)
        rgb = cond["rgb"][:, -self.img_cond_steps:]
        if self.num_img > 1:
            rgb = rgb.reshape(B, T_rgb, self.num_img, 3, H, W)
            rgb = einops.rearrange(rgb, "b t n c h w -> b n (t c) h w")
        elif self.num_img == 1:
            rgb = einops.rearrange(rgb, "b t c h w -> b (t c) h w")
        else:
            raise ValueError(f"self.num_img={self.num_img} < 1")
        rgb = rgb.float()
        if self.num_img == 2:
            rgb1, rgb2 = rgb[:, 0], rgb[:, 1]
            if self.augment:
                rgb1 = self.aug(rgb1)
                rgb2 = self.aug(rgb2)
            visual_feat1 = self.backbone.forward(rgb1)
            visual_feat1 = self.compress1.forward(visual_feat1, cond["state"].view(B, -1)) if hasattr(self, 'compress1') else self.compress(visual_feat1.flatten(1, -1))
            visual_feat2 = self.backbone.forward(rgb2)
            visual_feat2 = self.compress2.forward(visual_feat2, cond["state"].view(B, -1)) if hasattr(self, 'compress2') else self.compress(visual_feat2.flatten(1, -1))
            visual_feat = torch.cat([visual_feat1, visual_feat2], dim=-1)
        elif self.num_img == 1:
            if self.augment:
                rgb = self.aug(rgb)
            visual_feat = self.backbone.forward(rgb)
            if isinstance(self.compress, SpatialEmb):
                visual_feat = self.compress.forward(visual_feat, cond["state"].view(B, -1))
            else:
                visual_feat = self.compress(visual_feat.flatten(1, -1))
        else:
            raise NotImplementedError(f"num_img={self.num_img} not supported")
        
        # Combine visual and proprioceptive visual_features
        if self.embed_combination_type == 'add' or 'multiply':
            cond_emb = self.cond_embed(torch.cat([visual_feat, prop_emb], dim=-1))
        else:
            cond_emb = torch.cat([visual_feat, prop_emb], dim=-1)
        
        # Combine embeddings based on embed_combination_type
        if self.embed_combination_type == 'add':
            td_cond_emb = td_emb + cond_emb
        elif self.embed_combination_type == 'multiply':
            td_cond_emb = td_emb * cond_emb
        elif self.embed_combination_type == 'concate':
            td_cond_emb = torch.cat([td_emb, cond_emb], dim=-1)
        
        emd=torch.cat([action_embed, td_cond_emb], dim=-1)
        
        # Pass through Transformer blocks
        # omitted
        
        # Final layer to predict velocities
        velocity = self.vel_head(emd)
        if output_embedding:
            return velocity.view(B, Ta, Da), td_emb, cond_emb
        return velocity.view(B, Ta, Da)
    
    def sample_action(self,cond:dict,inference_steps:int,clip_intermediate_actions:bool,act_range:List[float], z:Tensor=None,save_chains:bool=False):
        """
        simply return action via integration (Euler's method). the initial noise could be specified. 
        when `save_chains` is True, also return the denoising trajectory.
        """
        B = cond['state'].shape[0]
        device=cond['state'].device

        x_hat:Tensor=z if z is not None else torch.randn(B, self.horizon_steps, self.action_dim, device=device)
        if save_chains:
            x_chain=torch.zeros((B, inference_steps+1, self.horizon_steps, self.action_dim), device=device)
        dt = (1 / inference_steps) * torch.ones_like(x_hat, device=device)
        steps = torch.linspace(0, 1-1/inference_steps, inference_steps, device=device).repeat(B, 1)
        for i in range(inference_steps):
            t = steps[:, i]
            dt_batch=(1 / inference_steps)* torch.ones(B, device=device)
            vt = self.forward(action=x_hat, time=t, dt=dt_batch, cond=cond, output_embedding=False)
            x_hat += vt * dt
            if clip_intermediate_actions or i == inference_steps-1: # always clip the output action. appended by Tonghe on 04/25/2025
                x_hat = x_hat.clamp(*act_range)
            if save_chains:
                x_chain[:, i+1] = x_hat
        if save_chains:
            return x_hat, x_chain
        return x_hat


class NoisyShortCutFlowMLP(NoisyFlowMLP):
    def __init__(
        self,
        policy:ShortCutFlowMLP,
        denoising_steps:int,
        learn_explore_noise_from:int,
        inital_noise_scheduler_type:str,
        min_logprob_denoising_std:float,
        max_logprob_denoising_std:float,
        learn_explore_time_embedding:bool,
        time_dim_explore:int,
        use_time_independent_noise:bool,
        device,
        noise_hidden_dims=None,
        activation_type='Tanh',
    ):  
        super().__init__(
            policy,
            denoising_steps,
            learn_explore_noise_from,
            inital_noise_scheduler_type,
            min_logprob_denoising_std,
            max_logprob_denoising_std,
            learn_explore_time_embedding,
            time_dim_explore,
            use_time_independent_noise,
            device,
            noise_hidden_dims,
            activation_type
        )
        self.policy:ShortCutFlowMLP
    
    # overload to receive shortcut features 
    def init_exploration_noise_net(self):
        if self.use_time_independent_noise:
            # sigma(s)
            # input dims for the noisy net
            noise_input_dim = self.policy.cond_enc_dim
            # hidden dims for the noisy net
            if not self.noise_hidden_dims:
                self.noise_hidden_dims = [16]
        else:
            if self.learn_explore_time_embedding:
                noise_input_dim = self.time_dim_explore + self.policy.cond_enc_dim
                self.time_embedding_explore = nn.Embedding(num_embeddings=self.denoising_steps, 
                                                       embedding_dim = self.time_dim_explore, 
                                                       device=self.device)
            else:
                # sigma(s,t)
                # input dims for the noisy net
                noise_input_dim = self.policy.td_emb_dim + self.policy.cond_enc_dim
                # hidden dims for the noisy net
                if not self.noise_hidden_dims:
                    self.noise_hidden_dims = [int(np.sqrt(noise_input_dim**2 + self.policy.act_dim_total**2))]
        
        self.explore_noise_net=ExploreNoiseNet(in_dim=noise_input_dim, 
                                                out_dim=self.policy.act_dim_total,
                                                logprob_denoising_std_range=[self.min_logprob_denoising_std, self.max_logprob_denoising_std], 
                                                device=self.device,
                                                hidden_dims=self.noise_hidden_dims,
                                                activation_type=self.noise_activation_type)

    # overload
    def forward(
        self,
        action,
        time,
        dt,
        cond,
        learn_exploration_noise=False,
        step=-1,
        verbose=False,
        **kwargs,
    )->Tuple[Tensor, Tensor]:
        """
        inputs:
            x: (B, Ta, Da)
            time: (B,) floating point in {0,1/2,1/4,1/8,...1/2^n} shortcut flow time
            cond: dict with key state/rgb; more recent obs at the end
                state: (B, To, Do)
            step: (B,) torch.tensor, optional, flow matching denoising step, from 0 to denoising_steps-1
            *here, B is the n_envs
        outputs:
             vel                [B, Ta, Da]
             noise_std          [B, Ta x Da]
        """
        B = action.shape[0]
        # WARNING: here you must secure that dt and time matches: time must be a multiple of 1.0 / self.denoising_steps.
        vel, td_emb, cond_emb = self.policy.forward(action, time, dt, cond, output_embedding=True)
        
        # noise head (for exploration). allow gradient flow.
        if self.initial_noise_scheduler_type=='const' or step < self.learn_explore_noise_from:
            noise_std       = self.logprob_noise_levels[:, step].repeat(B,1)
        else:
            if self.use_time_independent_noise:
                noise_feature    = cond_emb
            else:
                if self.learn_explore_time_embedding:
                    step_ts = torch.tensor(step, device = self.device).repeat(B)
                    time_emb_explore = self.time_embedding_explore(step_ts)
                    noise_feature    = torch.cat([time_emb_explore, cond_emb], dim=-1)
                else:
                    noise_feature    = torch.cat([td_emb.detach(), cond_emb], dim=-1)
            
            noise_std = self.explore_noise_net.forward(noise_feature=noise_feature)
            
            if verbose:
                log.info(f"step={step}, learnable noise = {noise_std.mean()}")
        if verbose:
            log.info(f"step={step}, set to learn from {self.learn_explore_noise_from}, will learn exploration noise ? {step >= self.learn_explore_noise_from}, noise_std={noise_std.mean()}require_grad={noise_std.requires_grad}")
        
        return vel, noise_std if learn_exploration_noise else noise_std.detach()

class NoisyVisionShortCutFlowMLP(NoisyShortCutFlowMLP):
    def __init__(
        self,
        policy:ShortCutFlowViT,
        denoising_steps:int,
        learn_explore_noise_from:int,
        inital_noise_scheduler_type:str,
        min_logprob_denoising_std:float,
        max_logprob_denoising_std:float,
        learn_explore_time_embedding:bool,
        time_dim_explore:int,
        use_time_independent_noise:bool,
        device,
        noise_hidden_dims=None,
        activation_type='Tanh',
    ):  
        super().__init__(
            policy,
            denoising_steps,
            learn_explore_noise_from,
            inital_noise_scheduler_type,
            min_logprob_denoising_std,
            max_logprob_denoising_std,
            learn_explore_time_embedding,
            time_dim_explore,
            use_time_independent_noise,
            device,
            noise_hidden_dims,
            activation_type,
        )
        self.policy:ShortCutFlowViT
    # overload
    def forward(
        self,
        action,
        time,
        cond,
        learn_exploration_noise=False,
        step=-1,
        verbose=False,
        **kwargs,
    )->Tuple[Tensor, Tensor]:
        """
        inputs:
            x: (B, Ta, Da)
            time: (B,) floating point in {0,1/2,1/4,1/8,...1/2^n} shortcut flow time
            cond: dict with key state/rgb; more recent obs at the end
                state: (B, To, Do)
            step: (B,) torch.tensor, optional, flow matching inference step, from 0 to denoising_steps-1
            *here, B is the n_envs
        outputs:
             vel                [B, Ta, Da]
             noise_std          [B, Ta x Da]
        """
        B = action.shape[0]
        # this is new for shortcut flows:
        dt = torch.full((B,), 1.0 / self.denoising_stepss, device=self.device)
        # WARNING: here you must secure that dt and time matches: time must be a multiple of 1.0 / self.denoising_steps.
        vel, td_emb, cond_emb = self.policy.forward(action, time, dt, cond, output_embedding=True)
        
        # noise head (for exploration). allow gradient flow.
        if self.initial_noise_scheduler_type=='const' or step < self.learn_explore_noise_from:
            noise_std       = self.logprob_noise_levels[:, step].repeat(B,1)
        else:
            if self.use_time_independent_noise:
                noise_feature    = cond_emb
            else:
                if self.learn_explore_time_embedding:
                    step_ts = torch.tensor(step, device = self.device).repeat(B)
                    time_emb_explore = self.time_embedding_explore(step_ts)
                    noise_feature    = torch.cat([time_emb_explore, cond_emb], dim=-1)
                else:
                    noise_feature    = torch.cat([td_emb.detach(), cond_emb], dim=-1)
            noise_std = self.explore_noise_net.forward(noise_feature=noise_feature)
        return vel, noise_std if learn_exploration_noise else noise_std.detach()