opentau 0.1.0__py3-none-any.whl

opentau/scripts/export_to_onnx.py

@@ -0,0 +1,180 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# 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 logging
+from pathlib import Path
+
+import torch
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.policies.factory import get_policy_class
+from opentau.policies.pi0.modeling_pi0 import PI0Policy
+from opentau.policies.pi05.modeling_pi05 import PI05Policy
+from opentau.utils.monkey_patch import (
+    torch_cumsum_patch,
+    torch_full_patch,
+    torch_pow_patch,
+)
+from opentau.utils.utils import auto_torch_device
+
+# Some patches are necessary only for dynamo export, which has current upstream bugs.
+# Nonetheless, we apply them here to ensure future compatibility.
+patches = [
+    torch_cumsum_patch,  # This is always necessary to load the ONNX artifact without error.
+    torch_full_patch,
+    torch_pow_patch,
+]
+
+KEY_STATES = "key_states"
+VALUE_STATES = "value_states"
+
+
+class InferenceWrapper(torch.nn.Module):
+    r"""Helper class to wrap the robot action decoder for ONNX export,
+    such that each input tensor is an individual argument to the `forward` method.
+    """
+
+    def __init__(
+        self,
+        decoder: PI0Policy | PI05Policy,
+        *,
+        prefix_pad_masks: torch.Tensor,
+        prefix_offsets: torch.Tensor,
+        num_cross_att_tokens: int,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.decoder = decoder
+        self.prefix_pad_masks = prefix_pad_masks
+        self.prefix_offsets = prefix_offsets
+        self.num_cross_att_tokens = num_cross_att_tokens
+        self.layer_idx = layer_idx
+
+    def forward(self, key_states, value_states, state):
+        vlm_tokens = (
+            {
+                self.layer_idx: {
+                    KEY_STATES: key_states,
+                    VALUE_STATES: value_states,
+                },
+            },
+            self.prefix_pad_masks,
+            self.prefix_offsets,
+            self.num_cross_att_tokens,
+        )
+        observation = {
+            "state": state,
+        }
+        actions = self.decoder.sample_actions(
+            observation,
+            vlm_token_cache_override=vlm_tokens,
+        )
+        return actions
+
+
+# Get the VLM cache for the dummy observation. This guarantees consistency with post-loading usage.
+def get_vlm_cache(cfg: TrainPipelineConfig, device: torch.device, dtype: torch.dtype):
+    logging.info("Getting VLM cache...")
+    policy_class = get_policy_class(cfg.policy.type)
+    cloud_vlm = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    cloud_vlm.set_execution_target("cloud")
+    cloud_vlm.to(device=device, dtype=torch.bfloat16)
+    cloud_vlm.eval()
+
+    vlm_camera_observation = {
+        f"camera{i}": torch.zeros((1, 3, *cfg.resolution), dtype=torch.bfloat16, device=device)
+        for i in range(cfg.num_cams)
+    }
+    vlm_observation = {
+        **vlm_camera_observation,
+        "prompt": ["Pick up yellow lego block and put it in the bin"],
+        "state": torch.zeros((1, cfg.max_state_dim), dtype=torch.bfloat16, device=device),
+        "img_is_pad": torch.zeros((1, cfg.num_cams), dtype=torch.bool, device=device),
+    }
+    cache, prefix_pad_masks, prefix_offsets, num_cross_att_tokens = cloud_vlm.get_vlm_tokens(vlm_observation)
+    assert len(cache) == 1, f"Expected only one cache entry for the dummy observation. Got {len(cache)}."
+    idx = list(cache)[0]
+    return (
+        cache[idx][KEY_STATES].to(dtype=dtype),
+        cache[idx][VALUE_STATES].to(dtype=dtype),
+        prefix_pad_masks,
+        prefix_offsets,
+        num_cross_att_tokens,
+        idx,
+    )
+
+
+@parser.wrap()
+def main(cfg: TrainPipelineConfig):
+    device = auto_torch_device()
+    dtype = torch.float32
+
+    # arguments for the dummy observation
+    (
+        key_states,
+        value_states,
+        prefix_pad_masks,
+        prefix_offsets,
+        num_cross_att_tokens,
+        layer_idx,
+    ) = get_vlm_cache(cfg, device, dtype)
+    state = torch.zeros((1, cfg.max_state_dim), device=device, dtype=dtype)
+    args = (key_states, value_states, state)
+    logging.info("Generated example args")
+
+    policy_class = get_policy_class(cfg.policy.type)
+    robot_action_decoder = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    robot_action_decoder.set_execution_target("robot")
+    robot_action_decoder.to(device)
+    robot_action_decoder.to(dtype=dtype)
+    robot_action_decoder.eval()
+    inference_wrapper = InferenceWrapper(
+        robot_action_decoder,
+        prefix_pad_masks=prefix_pad_masks,
+        prefix_offsets=prefix_offsets,
+        num_cross_att_tokens=num_cross_att_tokens,
+        layer_idx=layer_idx,
+    )
+    logging.info("Loaded policy")
+
+    logging.info("Applying monkey patches...")
+    for patch in patches:
+        patch()
+
+    logging.info("Exporting model to ONNX...")
+    with torch.inference_mode():
+        path = Path(cfg.policy.pretrained_path) / "robot_action_decoder.onnx"
+        path = path.resolve()
+        path.parent.mkdir(parents=True, exist_ok=True)  # Should be a no-op
+        print("Exporting model to ONNX at path:", path)
+        print("Current directory:", Path.cwd())
+        print("Trying to write to:", path)
+        with open(path, "wb"):
+            print("Write permissions check passed for:", path)
+        print("Running torch.onnx.export...")
+        torch.onnx.export(
+            inference_wrapper.eval(),
+            args,
+            path,
+            input_names=[KEY_STATES, VALUE_STATES, "state"],
+            output_names=["action_chunk"],
+            opset_version=18,
+            do_constant_folding=False,  # constant folding causes weird errors (getting dim -1 from a 0-dim scalar) after forward pass succeeds
+        )
+        logging.info(f"Successfully exported model to '{path}'.")
+
+
+if __name__ == "__main__":
+    main()

opentau/scripts/fake_tensor_training.py

@@ -0,0 +1,87 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# 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.
+
+from dataclasses import dataclass
+
+import torch
+
+from opentau.configs import parser
+from opentau.utils.fake_tensor import FakeTensorContext
+from opentau.utils.utils import auto_torch_device
+
+
+@dataclass
+class Args:
+    device: str | None = None
+    num_workers: int = 4
+    batch_size: int = 2
+    dim_in: int = 3
+    dim_out: int = 5
+    large_hidden_dim: int = 10**9
+
+
+@parser.wrap()
+def main(args: Args):
+    device = args.device or auto_torch_device()
+    data = [(torch.rand(args.dim_in), torch.rand(args.dim_out)) for _ in range(10)]
+    dataloader = torch.utils.data.DataLoader(
+        data,
+        num_workers=args.num_workers,
+        batch_size=args.batch_size,
+    )
+
+    # ---------------------------------------------------------------
+    # Everything inside this context will use FakeTensorMode
+    with FakeTensorContext():
+        # Create a model in FakeTensorContext shouldn't cost real memory for model parameters.
+        model = torch.nn.Sequential(
+            torch.nn.Linear(args.dim_in, args.large_hidden_dim),
+            torch.nn.Linear(args.large_hidden_dim, args.large_hidden_dim),
+            torch.nn.Linear(args.large_hidden_dim, args.dim_out),
+        ).to(device)
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
+    # End of FakeTensorContext
+    # ---------------------------------------------------------------
+
+    print("Model parameters: ")
+    for name, param in model.named_parameters():
+        print(f"{name}: {param})")
+    print()
+
+    losses = []
+    for i, (x, y) in enumerate(dataloader):
+        # ---------------------------------------------------------------
+        # Everything inside this context will use FakeTensorMode
+        # Ideally, we want to iterate the dataloader in FakeTensorMode as well.
+        # However, it does not work with multiple workers due to some serialization issue.
+        with FakeTensorContext():
+            x = x.to(device=device)
+            y = y.to(device=device)
+            optimizer.zero_grad()
+            loss = torch.nn.functional.mse_loss(model(x), y)
+            loss.backward()
+
+            optimizer.step()
+            losses.append(loss.item())
+            print(
+                f"Step {i}: symbolic loss = {loss.item()}, dummy numpy array = {loss.detach().cpu().numpy()}"
+            )
+        # End of FakeTensorContext
+        # ---------------------------------------------------------------
+
+    print("\nSymbolic mean loss:", sum(losses) / len(losses))
+
+
+if __name__ == "__main__":
+    main()

opentau/scripts/get_advantage_and_percentiles.py

@@ -0,0 +1,220 @@
+r"python src/opentau/scripts/get_advantage_and_percentiles.py  \
+--config_path=outputs/train/2025-11-29/00-38-59_value/checkpoints/00520000 \
+--batch_size=20 \
+--dataloader_batch_size=20 \
+--dataset_mixture=examples/advantage_config.json"
+
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# 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 json
+import logging
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+import draccus
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from opentau.configs import parser
+from opentau.configs.default import DatasetMixtureConfig
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.factory import make_dataset
+from opentau.datasets.lerobot_dataset import LeRobotDataset
+from opentau.datasets.utils import ADVANTAGES_PATH
+from opentau.policies.factory import get_policy_class
+from opentau.policies.value.reward import calculate_n_step_return
+from opentau.utils.random_utils import set_seed
+from opentau.utils.utils import (
+    auto_torch_device,
+    init_logging,
+)
+
+
+def ensure_primitive(maybe_tensor):
+    if isinstance(maybe_tensor, np.ndarray):
+        return ensure_primitive(torch.from_numpy(maybe_tensor))
+    if isinstance(maybe_tensor, torch.Tensor):
+        assert maybe_tensor.numel() == 1, f"Tensor must be a single value, got shape={maybe_tensor.numel()}"
+    return maybe_tensor
+
+
+_default0 = defaultdict(int)
+
+# Store dataset_mixture_path before filtering (needed for parsing inside main)
+# Handle both --dataset_mixture_path=<path> and --dataset_mixture=<path> (without nested fields)
+_dataset_mixture_path_value = None
+for arg in sys.argv:
+    if arg.startswith("--dataset_mixture_path="):
+        _dataset_mixture_path_value = arg.split("=", 1)[1]
+        break
+    elif arg.startswith("--dataset_mixture=") and "." not in arg.split("=", 1)[0]:
+        # --dataset_mixture=<path> without nested fields (e.g., not --dataset_mixture.datasets.0.repo_id=...)
+        _dataset_mixture_path_value = arg.split("=", 1)[1]
+        break
+
+# Create a wrapper that filters dataset_mixture path arguments before draccus parsing
+_original_wrap = parser.wrap()
+
+
+def _filter_dataset_mixture_path(fn):
+    """Wrapper that filters --dataset_mixture_path and --dataset_mixture=<path> from sys.argv before draccus sees it."""
+    wrapped_fn = _original_wrap(fn)
+
+    def filtered_wrapper(*args, **kwargs):
+        # If config is already provided, just call the function
+        if len(args) > 0:
+            return wrapped_fn(*args, **kwargs)
+
+        # Otherwise, filter dataset_mixture path arguments from sys.argv before draccus parses
+        original_argv = sys.argv.copy()
+        try:
+            filtered_args = []
+            for arg in sys.argv:
+                # Filter --dataset_mixture_path=<path>
+                if (
+                    arg.startswith("--dataset_mixture_path=")
+                    or arg.startswith("--dataset_mixture=")
+                    and "." not in arg.split("=", 1)[0]
+                ):
+                    continue
+                else:
+                    filtered_args.append(arg)
+            sys.argv = filtered_args
+            return wrapped_fn(*args, **kwargs)
+        finally:
+            sys.argv = original_argv
+
+    return filtered_wrapper
+
+
+@_filter_dataset_mixture_path
+def main(cfg: TrainPipelineConfig):
+    dataset_mixture_path = _dataset_mixture_path_value
+
+    if dataset_mixture_path:
+        logging.info(f"Loading dataset config from separate file: {dataset_mixture_path}")
+        mixture_cfg = draccus.parse(
+            config_class=DatasetMixtureConfig, config_path=dataset_mixture_path, args=[]
+        )
+    else:
+        logging.info("Using the dataset mixture config from the TrainPipelineConfig")
+        mixture_cfg = cfg.dataset_mixture
+
+    device = auto_torch_device()
+    # torch.autograd.set_detect_anomaly(True)
+
+    # TODO(shuheng): Do we need the random seed here?
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    logging.info("Creating policy")
+    policy_class = get_policy_class(cfg.policy.type)
+    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    policy.to(device=device, dtype=torch.bfloat16)
+    policy.eval()
+
+    # Advantages across all datasets
+    advantages = []
+
+    for dataset_idx, dataset_cfg in enumerate(mixture_cfg.datasets):
+        logging.info(f"Creating dataset {dataset_idx}")
+        dataset = make_dataset(dataset_cfg, cfg, return_advantage_input=True)
+        assert isinstance(dataset, LeRobotDataset), (
+            f"Expected instance of LeRobotDataset, got {type(dataset)}"
+        )
+        dataloader = DataLoader(
+            dataset,
+            batch_size=cfg.batch_size,
+            shuffle=False,
+            drop_last=False,
+            num_workers=cfg.num_workers,
+            pin_memory=torch.cuda.is_available(),
+            prefetch_factor=cfg.prefetch_factor,
+        )
+
+        values = {}
+        ds_advantage = {}  # per-dataset advantages
+        with torch.inference_mode():
+            # First pass to get the values
+            for batch in dataloader:
+                for key, value in batch.items():
+                    if isinstance(value, torch.Tensor):
+                        batch[key] = value.to(device)
+
+                for success, episode_index, episode_end_idx, current_idx, v0 in zip(
+                    batch["success"],
+                    batch["episode_index"],
+                    batch["episode_end_idx"],
+                    batch["current_idx"],
+                    policy.predict_value(batch),
+                    strict=True,
+                ):
+                    success, episode_index, episode_end_idx, current_idx, v0 = map(
+                        ensure_primitive, (success, episode_index, episode_end_idx, current_idx, v0)
+                    )
+                    reward = calculate_n_step_return(
+                        success=success,
+                        n_steps_look_ahead=cfg.policy.reward_config.N_steps_look_ahead,
+                        episode_end_idx=episode_end_idx,
+                        max_episode_length=cfg.policy.reward_config.reward_normalizer,
+                        current_idx=current_idx,
+                        c_neg=cfg.policy.reward_config.C_neg,
+                    )
+
+                    values[(episode_index, current_idx)] = {"v0": v0, "reward": reward}
+
+            # Second pass to compute the advantages
+            for batch in dataloader:
+                for episode_index, current_idx, timestamp in zip(
+                    batch["episode_index"],
+                    batch["current_idx"],
+                    batch["timestamp"],
+                    strict=True,
+                ):
+                    episode_index, current_idx, timestamp = map(
+                        ensure_primitive, (episode_index, current_idx, timestamp)
+                    )
+                    # check if the value for the next n_steps_look_ahead steps is available, else set it to 0
+                    look_ahead_idx = current_idx + cfg.policy.reward_config.N_steps_look_ahead
+                    vn = values.get((episode_index, look_ahead_idx), _default0)["v0"]
+                    reward = values.get((episode_index, current_idx), _default0)["reward"]
+                    v0 = values.get((episode_index, current_idx), _default0)["v0"]
+                    advantage = ensure_primitive(reward + vn - v0)
+                    advantages.append(advantage)
+                    ds_advantage[(episode_index, timestamp)] = advantage
+
+        # Convert tuple keys to strings for JSON serialization
+        advantage_data_json = {f"{ep_idx},{ts}": val for (ep_idx, ts), val in ds_advantage.items()}
+
+        # TODO(shuheng) avoid overwriting existing advantage files.
+        with open(Path(dataset.root) / ADVANTAGES_PATH, "w") as f:
+            json.dump(advantage_data_json, f, indent=4)
+
+    # Calculate percentiles of advantages: 0th, 5th, 10th, ..., 100th
+    percentiles = list(range(0, 101, 5))  # [0, 5, 10, 15, ..., 100]
+    advantage_percentiles = np.percentile(np.array(advantages), percentiles)
+
+    print("Advantage percentiles for deciding epsilon threshold:")
+    for p, val in zip(percentiles, advantage_percentiles, strict=False):
+        print(f"  {p:03d}th percentile: {val:.6f}")
+
+
+if __name__ == "__main__":
+    init_logging()
+    main()

opentau/scripts/high_level_planner_inference.py

@@ -0,0 +1,114 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env python
+
+import logging
+from dataclasses import asdict
+from pprint import pformat
+
+import torch
+from dotenv import load_dotenv
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.planner import HighLevelPlanner, Memory
+from opentau.policies.factory import get_policy_class
+from opentau.utils.random_utils import set_seed
+from opentau.utils.utils import (
+    init_logging,
+)
+
+load_dotenv()
+
+
+@parser.wrap()
+def inference_main(cfg: TrainPipelineConfig):
+    """
+    Reflects the whole pipeline from passing tasks to high level planner to generating actions from low level planner
+
+    Args:
+        cfg: configuration file. For example look at examples/dev_config.json
+    """
+
+    logging.info(pformat(asdict(cfg)))
+
+    # Check device is available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    logging.info("Creating policy")
+    policy_class = get_policy_class(cfg.policy.type)
+    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    policy.to(device)
+    policy.to(dtype=torch.bfloat16)
+    policy.eval()
+
+    hlp = HighLevelPlanner()
+    mem = Memory(len=2)
+
+    # compile the model if possible
+    if hasattr(torch, "compile"):
+        logging.info("Attempting to compile the policy with torch.compile()...")
+        try:
+            # Other options: "default", "max-autotune" (longer compile time)
+            policy = torch.compile(policy)
+            logging.info("Policy compiled successfully.")
+        except Exception as e:
+            logging.warning(f"torch.compile failed with error: {e}. Proceeding without compilation.")
+    else:
+        logging.warning(
+            "torch.compile is not available. Requires PyTorch 2.0+. Proceeding without compilation."
+        )
+
+    # Always reset policy before episode to clear out action cache.
+    policy.reset()
+
+    for i in range(5):
+        # create dummy observation for pi05
+        camera_observations = {
+            f"camera{i}": torch.zeros((1, 3, *cfg.resolution), dtype=torch.bfloat16, device=device)
+            for i in range(cfg.num_cams)
+        }
+
+        task = "Pick up yellow lego block and put it in the bin"
+
+        sub_task = hlp.inference(camera_observations, "", task, "gpt4o", mem).split("\n")[1].split('"')[1]
+
+        if mem:
+            mem.add_conversation("assistant", [{"type": "text", "text": sub_task}])
+
+        logging.info(f"{sub_task}")
+        observation = {
+            **camera_observations,
+            "state": torch.zeros((1, cfg.max_state_dim), dtype=torch.bfloat16, device=device),
+            "prompt": [sub_task],
+            "img_is_pad": torch.zeros((1, cfg.num_cams), dtype=torch.bool, device=device),
+            "action_is_pad": torch.zeros((1, cfg.action_chunk), dtype=torch.bool, device=device),
+        }
+
+        with torch.inference_mode():
+            for _ in range(1000):
+                action = policy.select_action(observation)
+                action = action.to("cpu").numpy()
+                print(f"Output dummy action: {action}")
+
+    logging.info("End of inference")
+
+
+if __name__ == "__main__":
+    init_logging()
+    inference_main()

opentau/scripts/inference.py

@@ -0,0 +1,70 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env python
+
+import logging
+from dataclasses import asdict
+from pprint import pformat
+
+import torch
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.policies.factory import get_policy_class
+from opentau.utils.random_utils import set_seed
+from opentau.utils.utils import (
+    attempt_torch_compile,
+    auto_torch_device,
+    create_dummy_observation,
+    init_logging,
+)
+
+
+@parser.wrap()
+def inference_main(cfg: TrainPipelineConfig):
+    logging.info(pformat(asdict(cfg)))
+
+    # Check device is available
+    device = auto_torch_device()
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    logging.info("Creating policy")
+    policy_class = get_policy_class(cfg.policy.type)
+    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    policy.to(device=device, dtype=torch.bfloat16)
+    policy.eval()
+    policy = attempt_torch_compile(policy, device_hint=device)
+
+    # Always reset policy before episode to clear out action cache.
+    policy.reset()
+
+    observation = create_dummy_observation(cfg, device, dtype=torch.bfloat16)
+
+    print(observation.keys())
+
+    with torch.inference_mode():
+        for _ in range(1000):
+            action = policy.select_action(observation)
+            action = action.to("cpu", torch.float32).numpy()
+            print(f"Output shape: {action.shape}")
+
+    logging.info("End of inference")
+
+
+if __name__ == "__main__":
+    init_logging()
+    inference_main()