actaug GR00T train+eval bundle: code (save-at/stop-at), eval stack (porting fixes), subset builder, protocol docs
38dc077 verified | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import os | |
| from functools import partial | |
| import torch | |
| import trt_torch as trt | |
| from transformers.feature_extraction_utils import BatchFeature | |
| def eagle_tensorrt_forward(self, vl_input): | |
| eagle_prefix = "eagle_" | |
| eagle_input = { | |
| k.removeprefix(eagle_prefix): v for k, v in vl_input.items() if k.startswith(eagle_prefix) | |
| } | |
| del eagle_input["image_sizes"] | |
| vl_input = eagle_input | |
| self.set_frozen_modules_to_eval_mode() | |
| batch_size = vl_input["pixel_values"].shape[0] | |
| position_ids = torch.arange(self.num_patches, device="cuda").expand((batch_size, -1)) | |
| if vl_input["pixel_values"].dtype != torch.float16: | |
| vl_input["pixel_values"] = vl_input["pixel_values"].to(torch.float16) | |
| assert ( | |
| vl_input["pixel_values"].shape[0] <= 8 | |
| ), "Batch size must be <= 8 because TensorRT engine was built with max_batch_size=8, you can try to adjust the max_batch_size in the build_engine.sh script and rebuild the engine." | |
| self.vit_engine.set_runtime_tensor_shape("pixel_values", vl_input["pixel_values"].shape) | |
| self.vit_engine.set_runtime_tensor_shape("position_ids", position_ids.shape) | |
| vit_embeds = self.vit_engine(vl_input["pixel_values"], position_ids)["vit_embeds"] | |
| self.llm_engine.set_runtime_tensor_shape("input_ids", vl_input["input_ids"].shape) | |
| self.llm_engine.set_runtime_tensor_shape("vit_embeds", vit_embeds.shape) | |
| self.llm_engine.set_runtime_tensor_shape("attention_mask", vl_input["attention_mask"].shape) | |
| embeddings = self.llm_engine(vl_input["input_ids"], vit_embeds, vl_input["attention_mask"])[ | |
| "embeddings" | |
| ] | |
| return BatchFeature( | |
| data={ | |
| "backbone_features": embeddings, | |
| "backbone_attention_mask": vl_input["attention_mask"], | |
| } | |
| ) | |
| def action_head_tensorrt_forward(self, backbone_output, action_input): | |
| # backbone_output = self.process_backbone_output(backbone_output) | |
| if backbone_output.backbone_features.dtype != torch.float16: | |
| backbone_output.backbone_features = backbone_output.backbone_features.to(torch.float16) | |
| self.vlln_vl_self_attention_engine.set_runtime_tensor_shape( | |
| "backbone_features", backbone_output.backbone_features.shape | |
| ) | |
| backbone_output.backbone_features = self.vlln_vl_self_attention_engine( | |
| backbone_output.backbone_features | |
| )["output"] | |
| vl_embeds = backbone_output.backbone_features | |
| embodiment_id = action_input.embodiment_id | |
| batch_size = vl_embeds.shape[0] | |
| if action_input.state.dtype != torch.float16: | |
| action_input.state = action_input.state.to(torch.float16) | |
| if embodiment_id.dtype != torch.int64: | |
| embodiment_id = embodiment_id.to(torch.int64) | |
| if vl_embeds.dtype != torch.float16: | |
| vl_embeds = vl_embeds.to(torch.float16) | |
| # Embed state with batch processing | |
| self.state_encoder_engine.set_runtime_tensor_shape("state", action_input.state.shape) | |
| self.state_encoder_engine.set_runtime_tensor_shape("embodiment_id", embodiment_id.shape) | |
| state_features = self.state_encoder_engine(action_input.state, embodiment_id)["output"] | |
| # Set initial actions as the sampled noise. | |
| device = vl_embeds.device | |
| actions = torch.randn( | |
| size=(batch_size, self.config.action_horizon, self.config.action_dim), | |
| dtype=vl_embeds.dtype, | |
| device=device, | |
| ) | |
| # This attribute is used to ensure the same actions is used for both PyTorch and TensorRT inference | |
| if hasattr(self, "init_actions"): | |
| actions = self.init_actions.expand((batch_size, -1, -1)) | |
| num_steps = self.num_inference_timesteps | |
| dt = 1.0 / num_steps | |
| for t in range(num_steps): | |
| t_cont = t / float(num_steps) # e.g. goes 0, 1/N, 2/N, ... | |
| t_discretized = int(t_cont * self.num_timestep_buckets) | |
| # Embed noised action trajectory with batch processing | |
| timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device) | |
| self.action_encoder_engine.set_runtime_tensor_shape("actions", actions.shape) | |
| self.action_encoder_engine.set_runtime_tensor_shape( | |
| "timesteps_tensor", timesteps_tensor.shape | |
| ) | |
| self.action_encoder_engine.set_runtime_tensor_shape("embodiment_id", embodiment_id.shape) | |
| action_features = self.action_encoder_engine(actions, timesteps_tensor, embodiment_id)[ | |
| "output" | |
| ] | |
| # Maybe add position embedding. | |
| if self.config.add_pos_embed: | |
| pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) | |
| pos_embs = self.position_embedding(pos_ids).unsqueeze(0).to(torch.float16) | |
| action_features = action_features + pos_embs | |
| vl_embs = vl_embeds | |
| # Join vision, language, state and action embedding along sequence dimension. | |
| sa_embs = torch.cat((state_features, action_features), dim=1) | |
| # Run model forward with batch processing | |
| if vl_embs.dtype != torch.float16: | |
| vl_embs = vl_embs.to(torch.float16) | |
| self.DiT_engine.set_runtime_tensor_shape("vl_embs", vl_embs.shape) | |
| self.DiT_engine.set_runtime_tensor_shape("sa_embs", sa_embs.shape) | |
| self.DiT_engine.set_runtime_tensor_shape("timesteps_tensor", timesteps_tensor.shape) | |
| model_output = self.DiT_engine(sa_embs, vl_embs, timesteps_tensor)["output"] | |
| self.action_decoder_engine.set_runtime_tensor_shape("model_output", model_output.shape) | |
| self.action_decoder_engine.set_runtime_tensor_shape("embodiment_id", embodiment_id.shape) | |
| pred = self.action_decoder_engine(model_output, embodiment_id)["output"] | |
| pred_velocity = pred[:, -self.action_horizon :] | |
| # Update actions using euler integration. | |
| actions = actions + dt * pred_velocity | |
| return BatchFeature(data={"action_pred": actions}) | |
| def setup_tensorrt_engines(policy, trt_engine_path): | |
| """ | |
| Setup TensorRT engines for GR00T model inference. | |
| Args: | |
| policy: GR00T policy model instance | |
| """ | |
| policy.model.backbone.num_patches = ( | |
| policy.model.backbone.eagle_model.vision_model.vision_model.embeddings.num_patches | |
| ) | |
| if hasattr(policy.model.backbone.eagle_model, "vision_model"): | |
| del policy.model.backbone.eagle_model.vision_model | |
| if hasattr(policy.model.backbone.eagle_model, "language_model"): | |
| del policy.model.backbone.eagle_model.language_model | |
| if hasattr(policy.model.action_head, "vlln"): | |
| del policy.model.action_head.vlln | |
| if hasattr(policy.model.action_head, "vl_self_attention"): | |
| del policy.model.action_head.vl_self_attention | |
| if hasattr(policy.model.action_head, "model"): | |
| del policy.model.action_head.model | |
| if hasattr(policy.model.action_head, "state_encoder"): | |
| del policy.model.action_head.state_encoder | |
| if hasattr(policy.model.action_head, "action_encoder"): | |
| del policy.model.action_head.action_encoder | |
| if hasattr(policy.model.action_head, "action_decoder"): | |
| del policy.model.action_head.action_decoder | |
| torch.cuda.empty_cache() | |
| # Setup backbone engines | |
| policy.model.backbone.vit_engine = trt.Engine(os.path.join(trt_engine_path, "vit.engine")) | |
| policy.model.backbone.llm_engine = trt.Engine(os.path.join(trt_engine_path, "llm.engine")) | |
| # Setup action head engines | |
| policy.model.action_head.vlln_vl_self_attention_engine = trt.Engine( | |
| os.path.join(trt_engine_path, "vlln_vl_self_attention.engine") | |
| ) | |
| policy.model.action_head.action_encoder_engine = trt.Engine( | |
| os.path.join(trt_engine_path, "action_encoder.engine") | |
| ) | |
| policy.model.action_head.action_decoder_engine = trt.Engine( | |
| os.path.join(trt_engine_path, "action_decoder.engine") | |
| ) | |
| policy.model.action_head.DiT_engine = trt.Engine(os.path.join(trt_engine_path, "DiT.engine")) | |
| policy.model.action_head.state_encoder_engine = trt.Engine( | |
| os.path.join(trt_engine_path, "state_encoder.engine") | |
| ) | |
| # Set TensorRT forward functions | |
| policy.model.backbone.forward = partial(eagle_tensorrt_forward, policy.model.backbone) | |
| policy.model.action_head.get_action = partial( | |
| action_head_tensorrt_forward, policy.model.action_head | |
| ) | |