| { | |
| "__config_type__": "robo_orchard_lab.models.bip3d.structure:BIP3DConfig", | |
| "class_type": "robo_orchard_lab.models.bip3d.structure:BIP3D", | |
| "backbone": { | |
| "type": "robo_orchard_lab.models.modules.swin_transformer:SwinTransformer", | |
| "embed_dims": 96, | |
| "depths": [ | |
| 2, | |
| 2, | |
| 6, | |
| 2 | |
| ], | |
| "num_heads": [ | |
| 3, | |
| 6, | |
| 12, | |
| 24 | |
| ], | |
| "window_size": 7, | |
| "mlp_ratio": 4, | |
| "qkv_bias": true, | |
| "qk_scale": null, | |
| "drop_rate": 0.0, | |
| "attn_drop_rate": 0.0, | |
| "out_indices": [ | |
| 1, | |
| 2, | |
| 3 | |
| ], | |
| "with_cp": true, | |
| "convert_weights": false | |
| }, | |
| "decoder": { | |
| "type": "robo_orchard_lab.models.holobrain.action_decoder:HoloBrainActionDecoder", | |
| "head": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:UpsampleHead", | |
| "upsample_sizes": [ | |
| 32, | |
| 64 | |
| ], | |
| "input_dim": 256, | |
| "dims": [ | |
| 128, | |
| 64 | |
| ], | |
| "norm": { | |
| "type": "torch.nn.modules.normalization:RMSNorm", | |
| "normalized_shape": 256 | |
| }, | |
| "act": { | |
| "type": "torch.nn.modules.activation:SiLU", | |
| "inplace": true | |
| }, | |
| "norm_act_idx": [ | |
| 0, | |
| 1 | |
| ], | |
| "num_output_layers": 2, | |
| "out_dim": 8 | |
| }, | |
| "mobile_head": null, | |
| "transformer_cfg": { | |
| "__config_type__": "robo_orchard_lab.models.holobrain.action_decoder:HoloBrainDecoderTransformerConfig", | |
| "img_cross_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:RotaryAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8, | |
| "max_position_embeddings": 32 | |
| }, | |
| "norm_layer": { | |
| "type": "torch.nn.modules.normalization:RMSNorm", | |
| "normalized_shape": 256 | |
| }, | |
| "ffn": { | |
| "type": "robo_orchard_lab.models.layers.transformer_layers:FFN", | |
| "embed_dims": 256, | |
| "feedforward_channels": 2048, | |
| "act_cfg": { | |
| "type": "torch.nn.modules.activation:SiLU", | |
| "inplace": true | |
| } | |
| }, | |
| "operation_order": [ | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp", | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp", | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp", | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp", | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp", | |
| "t_norm", | |
| "temp_joint_attn", | |
| "gate_msa", | |
| "norm", | |
| "img_cross_attn", | |
| "norm", | |
| "text_cross_attn", | |
| "norm", | |
| "scale_shift", | |
| "ffn", | |
| "gate_mlp" | |
| ], | |
| "joint_self_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:JointGraphAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8 | |
| }, | |
| "temp_cross_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:RotaryAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8, | |
| "max_position_embeddings": 32 | |
| }, | |
| "text_cross_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:RotaryAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8, | |
| "max_position_embeddings": 512 | |
| }, | |
| "temp_joint_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:TemporalJointGraphAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8, | |
| "max_position_embeddings": 32 | |
| }, | |
| "timestep_norm_layer": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:AdaRMSNorm", | |
| "normalized_shape": 256, | |
| "condition_dims": 256, | |
| "zero": true | |
| }, | |
| "pre_norm": true | |
| }, | |
| "base_cfg": { | |
| "__config_type__": "robo_orchard_lab.models.holobrain.action_decoder:HoloBrainDecoderBaseConfig", | |
| "training_noise_scheduler": { | |
| "type": "diffusers.schedulers.scheduling_ddpm:DDPMScheduler", | |
| "num_train_timesteps": 1000, | |
| "beta_schedule": "squaredcos_cap_v2", | |
| "prediction_type": "sample", | |
| "clip_sample": false | |
| }, | |
| "test_noise_scheduler": { | |
| "type": "diffusers.schedulers.scheduling_dpmsolver_multistep:DPMSolverMultistepScheduler", | |
| "num_train_timesteps": 1000, | |
| "beta_schedule": "squaredcos_cap_v2", | |
| "prediction_type": "sample" | |
| }, | |
| "num_inference_timesteps": 10, | |
| "feature_level": [ | |
| 1, | |
| 2 | |
| ], | |
| "state_dims": 8, | |
| "embed_dims": 256, | |
| "pred_steps": 64, | |
| "act_cfg": { | |
| "type": "torch.nn.modules.activation:SiLU", | |
| "inplace": true | |
| }, | |
| "num_test_traj": 1, | |
| "chunk_size": 4, | |
| "force_kinematics": false, | |
| "with_mobile": false, | |
| "mobile_traj_state_dims": 2, | |
| "use_joint_mask": true, | |
| "noise_type": "local_joint", | |
| "pred_scaled_joint": false, | |
| "prediction_type": "relative_joint_relative_pose" | |
| }, | |
| "training_cfg": { | |
| "__config_type__": "robo_orchard_lab.models.holobrain.action_decoder:HoloBrainTrainingConfig", | |
| "loss": { | |
| "type": "robo_orchard_lab.models.holobrain.loss:HoloBrainActionLoss", | |
| "timestep_loss_weight": 1000, | |
| "parallel_loss_weight": 0.1, | |
| "smooth_l1_beta": 0.04, | |
| "loss_mode": "smooth_l1" | |
| }, | |
| "temporal_attn_drop": 0.05, | |
| "num_parallel_training_sample": 4, | |
| "teacher_forcing_rate": 0.02, | |
| "teacher_forcing_mean_steps": null | |
| }, | |
| "robot_encoder": { | |
| "type": "robo_orchard_lab.models.holobrain.robot_state_encoder:HoloBrainRobotStateEncoder", | |
| "transformer_cfg": { | |
| "__config_type__": "robo_orchard_lab.models.holobrain.robot_state_encoder:HoloBrainEncoderTransformerConfig", | |
| "joint_self_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:JointGraphAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8 | |
| }, | |
| "norm_layer": { | |
| "type": "torch.nn.modules.normalization:RMSNorm", | |
| "normalized_shape": 256 | |
| }, | |
| "ffn": { | |
| "type": "robo_orchard_lab.models.layers.transformer_layers:FFN", | |
| "embed_dims": 256, | |
| "feedforward_channels": 2048, | |
| "act_cfg": { | |
| "type": "torch.nn.modules.activation:SiLU", | |
| "inplace": true | |
| } | |
| }, | |
| "operation_order": [ | |
| "norm", | |
| "joint_self_attn", | |
| null, | |
| null, | |
| "norm", | |
| "ffn", | |
| "norm", | |
| "joint_self_attn", | |
| null, | |
| null, | |
| "norm", | |
| "ffn", | |
| "norm", | |
| "joint_self_attn", | |
| null, | |
| null, | |
| "norm", | |
| "ffn", | |
| "norm", | |
| "joint_self_attn", | |
| null, | |
| null, | |
| "norm", | |
| "ffn", | |
| "norm" | |
| ], | |
| "temp_self_attn": { | |
| "type": "robo_orchard_lab.models.holobrain.layers:RotaryAttention", | |
| "embed_dims": 256, | |
| "num_heads": 8, | |
| "max_position_embeddings": 32 | |
| }, | |
| "pre_norm": true | |
| }, | |
| "base_cfg": { | |
| "__config_type__": "robo_orchard_lab.models.holobrain.robot_state_encoder:HoloBrainEncoderBaseConfig", | |
| "embed_dims": 256, | |
| "state_dims": 8, | |
| "act_cfg": { | |
| "type": "torch.nn.modules.activation:SiLU", | |
| "inplace": true | |
| }, | |
| "chunk_size": 1 | |
| } | |
| } | |
| }, | |
| "neck": { | |
| "type": "robo_orchard_lab.models.modules.channel_mapper:ChannelMapper", | |
| "in_channels": [ | |
| 192, | |
| 384, | |
| 768 | |
| ], | |
| "kernel_size": 1, | |
| "out_channels": 256, | |
| "act_cfg": null, | |
| "bias": true, | |
| "norm_cfg": { | |
| "type": "torch.nn.modules.normalization:GroupNorm", | |
| "num_groups": 32 | |
| }, | |
| "num_outs": 4 | |
| }, | |
| "text_encoder": { | |
| "type": "robo_orchard_lab.models.bip3d.bert:BertModel", | |
| "special_tokens_list": [ | |
| "[CLS]", | |
| "[SEP]" | |
| ], | |
| "name": "./ckpt/bert-base-uncased", | |
| "pad_to_max": false, | |
| "use_sub_sentence_represent": true, | |
| "add_pooling_layer": false, | |
| "max_tokens": 768, | |
| "use_checkpoint": true, | |
| "return_tokenized": true | |
| }, | |
| "feature_enhancer": { | |
| "type": "robo_orchard_lab.models.bip3d.feature_enhancer:TextImageDeformable2DEnhancer", | |
| "embed_dims": 256, | |
| "num_layers": 6, | |
| "text_img_attn_block": { | |
| "v_dim": 256, | |
| "l_dim": 256, | |
| "embed_dim": 1024, | |
| "num_heads": 4, | |
| "init_values": 0.0001 | |
| }, | |
| "img_attn_block": { | |
| "self_attn_cfg": { | |
| "embed_dims": 256, | |
| "num_levels": 4, | |
| "im2col_step": 1 | |
| }, | |
| "ffn_cfg": { | |
| "embed_dims": 256, | |
| "feedforward_channels": 2048, | |
| "ffn_drop": 0.0 | |
| } | |
| }, | |
| "text_attn_block": { | |
| "self_attn_cfg": { | |
| "num_heads": 4, | |
| "embed_dims": 256 | |
| }, | |
| "ffn_cfg": { | |
| "embed_dims": 256, | |
| "feedforward_channels": 1024, | |
| "ffn_drop": 0.0 | |
| } | |
| }, | |
| "num_feature_levels": 4, | |
| "positional_encoding": { | |
| "num_feats": 128, | |
| "normalize": true, | |
| "offset": 0.0, | |
| "temperature": 20 | |
| } | |
| }, | |
| "spatial_enhancer": { | |
| "type": "robo_orchard_lab.models.bip3d.spatial_enhancer:DepthFusionSpatialEnhancer", | |
| "embed_dims": 256, | |
| "feature_3d_dim": 128, | |
| "num_depth_layers": 2, | |
| "min_depth": 0.01, | |
| "max_depth": 1.2, | |
| "num_depth": 128, | |
| "with_feature_3d": true, | |
| "loss_depth_weight": 1.0 | |
| }, | |
| "data_preprocessor": { | |
| "type": "robo_orchard_lab.models.layers.data_preprocessors:BaseDataPreprocessor", | |
| "channel_flip": true, | |
| "unsqueeze_depth_channel": true, | |
| "mean": [ | |
| 123.675, | |
| 116.28, | |
| 103.53 | |
| ], | |
| "std": [ | |
| 58.395, | |
| 57.12, | |
| 57.375 | |
| ], | |
| "batch_transforms": [ | |
| { | |
| "type": "robo_orchard_lab.models.bip3d.spatial_enhancer:BatchDepthProbGTGenerator", | |
| "min_depth": 0.01, | |
| "max_depth": 1.2, | |
| "num_depth": 128, | |
| "origin_stride": 2, | |
| "valid_threshold": 0.5, | |
| "stride": [ | |
| 8, | |
| 16, | |
| 32, | |
| 64 | |
| ] | |
| } | |
| ] | |
| }, | |
| "backbone_3d": { | |
| "type": "robo_orchard_lab.models.modules.swin_transformer:SwinTransformer", | |
| "in_channels": 1, | |
| "embed_dims": 16, | |
| "depths": [ | |
| 2, | |
| 2, | |
| 6, | |
| 2 | |
| ], | |
| "num_heads": [ | |
| 4, | |
| 8, | |
| 8, | |
| 16 | |
| ], | |
| "window_size": 7, | |
| "mlp_ratio": 4, | |
| "qkv_bias": true, | |
| "qk_scale": null, | |
| "drop_rate": 0.0, | |
| "attn_drop_rate": 0.0, | |
| "out_indices": [ | |
| 1, | |
| 2, | |
| 3 | |
| ], | |
| "with_cp": true, | |
| "convert_weights": false | |
| }, | |
| "neck_3d": { | |
| "type": "robo_orchard_lab.models.modules.channel_mapper:ChannelMapper", | |
| "in_channels": [ | |
| 32, | |
| 64, | |
| 128 | |
| ], | |
| "kernel_size": 1, | |
| "out_channels": 128, | |
| "act_cfg": null, | |
| "bias": true, | |
| "norm_cfg": { | |
| "type": "torch.nn.modules.normalization:GroupNorm", | |
| "num_groups": 4 | |
| }, | |
| "num_outs": 4 | |
| }, | |
| "input_2d": "imgs", | |
| "input_3d": "depths", | |
| "embed_dims": 256, | |
| "pre_spatial_enhancer": false | |
| } |