opentau 0.1.0__py3-none-any.whl

opentau/planner/high_level_planner.py

@@ -0,0 +1,366 @@
+# 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 abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+from openai import OpenAI
+from PIL import Image
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForVision2Seq,
+    AutoProcessor,
+    LlamaTokenizer,
+)
+
+from opentau.planner.utils.memory import Memory
+from opentau.planner.utils.utils import load_prompt_library, tensor_to_base64
+
+
+class BaseHighLevelPlanner(ABC):
+    """
+    Represents a High level planner which has ability to infer various open source models and closed models like gpt-4o.
+    Generates a list of low level plans given a high level plan
+    """
+
+    def __init__(self):
+        self.prompts_dict = load_prompt_library("src/opentau/prompts/planner/prompts.yaml")
+
+    @abstractmethod
+    def inference(
+        self, image_dict: dict[str, torch.Tensor], model_name: str, task: str, mem: Optional[Memory] = None
+    ) -> str:
+        """
+        Handles inferencing the planner given images and language inputs
+        """
+
+        pass
+
+    def calculate_usage(self, response) -> float:
+        """
+        Calculates cost for each call to gpt-4o
+
+        Args:
+
+            response : a response object from gpt chat compeletion method
+
+        Returns:
+
+            cost (float) : cost for one call
+        """
+
+        prompt_tokens = response.usage.prompt_tokens
+        completion_tokens = response.usage.completion_tokens
+
+        cost = (prompt_tokens / 1000) * 0.0025 + (completion_tokens / 1000) * 0.01
+
+        return cost
+
+
+class HighLevelPlanner(BaseHighLevelPlanner):
+    """
+    Represents a High level planner which has ability to infer various open source models and closed models like gpt-4o.
+    Generates a list of low level plans given a high level plan
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def model_and_tokenizer(self, model_name: str, device: str):
+        if model_name == "cogvlm-chat-hf":
+            processor = LlamaTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")
+            model = AutoModelForCausalLM.from_pretrained(
+                "THUDM/cogvlm-chat-hf",
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+            ).eval()
+        elif model_name == "SmolVLM-256M-Instruct":
+            processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-256M-Instruct")
+            model = AutoModelForVision2Seq.from_pretrained(
+                "HuggingFaceTB/SmolVLM-256M-Instruct", torch_dtype=torch.float16
+            ).to(device)
+        elif model_name == "SmolVLM-500M-Instruct":
+            processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-500M-Instruct")
+            model = AutoModelForVision2Seq.from_pretrained(
+                "HuggingFaceTB/SmolVLM-500M-Instruct", torch_dtype=torch.float16
+            ).to(device)
+        elif model_name == "SmolVLM2-2.2B-Instruct":
+            processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct")
+            model = AutoModelForImageTextToText.from_pretrained(
+                "HuggingFaceTB/SmolVLM2-2.2B-Instruct", torch_dtype=torch.float32
+            ).to(device)
+        else:
+            raise RuntimeError(f"The specified model  {model_name} is not supported")
+
+        return model, processor
+
+    def generate_prompt(self, task: str, mem: Memory | None) -> str:
+        """
+        Generates prompt for gpt-4o model based on memory.
+
+        Args:
+
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+
+        Returns:
+
+            prompt (str): a prompt
+        """
+
+        if mem is None:
+            prompt = (
+                f"Look at the image <image>. The task is {task}."
+                + "\n"
+                + self.prompts_dict["prompts"]["robot_user_without_memory_manipulation"]["template"]
+            )
+
+        else:
+            prompt = (
+                f"Look at the image <image>. The task is {task}."
+                + "\n"
+                + self.prompts_dict["prompts"]["robot_user_with_memory_manipulation"]["template"]
+            )
+
+        return prompt
+
+    def gpt_inference(self, image_dict: dict[str, torch.Tensor], task: str, mem: Memory | None) -> str:
+        """
+        Calls openai Api and passes high level plan and memory
+
+        Args:
+
+            image_dict(dict[str , torch.Tensor]) : dict of tensors of images in base64 format
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+
+        Returns:
+
+            response (str): a low level language command that can be understood by low level planner
+        """
+
+        images = tensor_to_base64(image_dict, "mani")
+
+        client = OpenAI()
+
+        prompt = self.generate_prompt(task, mem)
+
+        content = [
+            {"type": "text", "text": f"{prompt}"},
+        ]
+
+        for image_base64 in images:
+            content.append(
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}}
+            )
+
+        if mem is None:
+            message = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"{self.prompts_dict['prompts']['robot_system_manipulation']['template']}",
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": content,
+                },
+            ]
+        else:
+            mem.add_conversation("user", content)
+            message = mem.get_conversation()
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=message,
+            max_tokens=500,
+            temperature=0.0,
+            tool_choice=None,
+        )
+
+        res = response.choices[0].message.content
+
+        return res
+
+    def opensource_inference(self, image_path, task, model_name):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        model, processor = self.model_and_tokenizer(model_name, device)
+
+        prompt = self.generate_prompt(task)
+
+        image1 = Image.open(image_path).convert("RGB")
+
+        inputs = processor(text=prompt, images=[image1], return_tensors="pt")
+        inputs = inputs.to(device)
+
+        # Generate outputs
+        generated_ids = model.generate(**inputs, max_new_tokens=500)
+        generated_texts = processor.batch_decode(
+            generated_ids,
+            skip_special_tokens=True,
+        )
+
+        return generated_texts[0]
+
+    def inference(
+        self, image_dict: dict[str, torch.Tensor], model_name: str, task: str, mem: Optional[Memory] = None
+    ) -> str:
+        """
+        Handles calling of open source models and gpt-4o models
+
+        Args:
+
+            image_dict (dict[str , torch.Tensor]) : dict of tensors of images in base64 format
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+            model_name (str) : Name of open source model to be inferenced. To use gpt models, pass a gpt4o
+            image_path (str) : Path to image file for opensource models.
+
+        Returns:
+
+            response (str): a low level language command that can be understood by low level planner
+        """
+
+        if model_name == "gpt4o":
+            actions = self.gpt_inference(image_dict, task, mem)
+            actions = actions.split("```json")[1].split("[")[1].split("]")[0]
+        else:
+            actions = self.opensource_inference(model_name, task)
+
+        return actions
+
+
+class NavHighLevelPlanner(BaseHighLevelPlanner):
+    """
+    Represents a High level planner which has ability to infer various open source models and closed models like gpt-4o.
+    Generates a list of low level plans given a high level plan
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def generate_prompt(self, task: str, mem: Memory | None) -> str:
+        """
+        Generates prompt for gpt-4o model based on memory.
+
+        Args:
+
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+
+        Returns:
+
+            prompt (str): a prompt
+        """
+
+        prompt = (
+            self.prompts_dict["prompts"]["robot_user_navigation"]["template"]
+            + f"Look at the given images {' '.join(['<image>'] * 21)} from starting point. The task is {task}."
+        )
+
+        return prompt
+
+    def gpt_inference(self, image_dict: dict[str, torch.Tensor], task: str, mem: Memory | None) -> str:
+        """
+        Calls openai Api and passes high level plan and memory
+
+        Args:
+
+            image_dict(dict[str , torch.Tensor]) : dict of tensors of images in base64 format
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+
+        Returns:
+
+            response (str): a low level language command that can be understood by low level planner
+        """
+
+        images = tensor_to_base64(image_dict, "nav")
+
+        client = OpenAI()
+
+        prompt = self.generate_prompt(task, mem)
+
+        content = [
+            {"type": "text", "text": f"{prompt}"},
+        ]
+
+        for image_base64 in images:
+            content.append(
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}}
+            )
+
+        if mem is None:
+            message = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": f"{self.prompts_dict['prompts']['robot_system_navigation']['template']}",
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": content,
+                },
+            ]
+        else:
+            mem.add_conversation("user", content)
+            message = mem.get_conversation()
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=message,
+            max_tokens=500,
+            temperature=0.0,
+            tool_choice=None,
+        )
+
+        res = response.choices[0].message.content
+
+        return res
+
+    def inference(
+        self, image_dict: dict[str, torch.Tensor], model_name: str, task: str, mem: Optional[Memory] = None
+    ) -> str:
+        """
+        Handles calling of open source models and gpt-4o models
+
+        Args:
+
+            image_dict (dict[str , torch.Tensor]) : dict of tensors of images in base64 format
+            task (str): a high level command as a language
+            mem (Memory|None): instance of Memory class from utils file, which stores all the previous conservations of relevant task in hand. Its set to None if memory is no required while inferencing.
+            model_name (str) : Name of open source model to be inferenced. To use gpt models, pass a gpt4o
+            image_path (str) : Path to image file for opensource models.
+
+        Returns:
+
+            response (str): a low level language command that can be understood by low level planner
+        """
+
+        if model_name == "gpt4o":
+            actions = self.gpt_inference(image_dict, task, mem)
+
+        return actions

opentau/planner/utils/memory.py

@@ -0,0 +1,64 @@
+# 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 collections import deque
+
+from opentau.planner.utils.utils import load_prompt_library
+
+
+class Memory:
+    """
+    Generates a memory like class to store, retrieve all the conversations for a particular task between user and LLM assistant
+    """
+
+    def __init__(self, conversation=None, len=1000):
+        """
+        Initializes conservation variable if any earlier conservation needs to be append before initilaizing the memory object.
+        len: fixed buffer size of conversation
+        """
+
+        self.len = len
+        self.prompts_dict = load_prompt_library("src/opentau/planner/prompts.yaml")
+
+        if conversation:
+            self.conversation = conversation
+        else:
+            self.conversation = deque()
+            context = [
+                {
+                    "type": "text",
+                    "text": f"{self.prompts_dict['prompts']['robot_system_manipulation']['template']}",
+                }
+            ]
+            self.add_conversation("system", context)
+
+    def add_conversation(self, role: str, message: list[dict[str, str]]) -> None:
+        """
+        Adds new conversations to history of conversations.
+
+        Args:
+            role (str): The message given by (system, user, assistant)
+            message (list[dict[str, str]]): message containing text and/or images.
+        """
+
+        if len(self.conversation) >= self.len:
+            self.conversation.popleft()
+
+        self.conversation.append({"role": role, "content": message})
+
+    def get_conversation(self) -> list[dict[str, str]]:
+        """
+        Returns the stored conversation or history of conversations
+        """
+        return list(self.conversation)

opentau/planner/utils/utils.py

@@ -0,0 +1,65 @@
+# 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 base64
+import io
+
+import torch
+import yaml
+from PIL import Image
+
+
+def tensor_to_base64(image_dict: dict[str, torch.Tensor], task: str) -> list[bytes]:
+    """
+    converts dictionary of tensors into list of base64
+
+    Args :
+        image_dict (dict[str, torch.Tensor]) : Dictionary of tensors (camera images)
+        task : two supported tasks, i.e , 'mani' and 'nav'
+
+    Return:
+        images (list[bytes]) : List of base64
+    """
+
+    images = []
+
+    for img in image_dict.values():
+        if task == "mani":
+            img_tensor = img.squeeze(0)
+            img_tensor = img_tensor.to(dtype=torch.float32, device="cpu")
+            img_tensor = img_tensor.clamp(0, 1) * 255.0
+            img = Image.fromarray(img_tensor.to(torch.uint8).permute(1, 2, 0).numpy())
+
+        buffered = io.BytesIO()
+        img.save(buffered, format="PNG")
+
+        images.append(base64.b64encode(buffered.getvalue()).decode("utf-8"))
+
+    return images
+
+
+def load_prompt_library(filepath: str) -> dict:
+    """
+    Loads a YAML file and returns its content as a dictionary.
+    """
+    try:
+        with open(filepath) as file:
+            # Use yaml.safe_load() to parse the YAML file
+            return yaml.safe_load(file)
+    except FileNotFoundError:
+        print(f"Error: The file at '{filepath}' was not found.")
+        return None
+    except yaml.YAMLError as e:
+        print(f"Error parsing the YAML file: {e}")
+        return None

opentau/policies/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+
+"""Policies module for OpenTau.
+
+This module exports the configuration classes for available policies,
+such as PI0, PI05, and Value policy.
+"""
+
+from .pi0.configuration_pi0 import PI0Config as PI0Config
+from .pi05.configuration_pi05 import PI05Config as PI05Config
+from .value.configuration_value import ValueConfig as ValueConfig

opentau/policies/factory.py

@@ -0,0 +1,172 @@
+#!/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.
+
+"""Factory functions for creating policy instances and configurations.
+
+This module provides utility functions to instantiate policy classes and their
+corresponding configurations based on policy names and types. It handles the
+logic for creating fresh policies or loading pretrained ones, as well as
+parsing features from datasets or environments to properly configure the policies.
+"""
+
+from typing import Optional
+
+import numpy as np
+from torch import nn
+
+from opentau.configs.policies import PreTrainedConfig
+from opentau.configs.types import FeatureType
+from opentau.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from opentau.datasets.utils import dataset_to_policy_features
+from opentau.policies.pi0.configuration_pi0 import PI0Config
+from opentau.policies.pi05.configuration_pi05 import PI05Config
+from opentau.policies.pretrained import PreTrainedPolicy
+from opentau.policies.value.configuration_value import ValueConfig
+
+
+def get_policy_class(name: str) -> type[PreTrainedPolicy]:
+    """Get the policy's class given a name.
+
+    Args:
+        name: The name of the policy (e.g., "pi0", "pi05", "value").
+            Must match the policy class's `name` attribute.
+
+    Returns:
+        type[PreTrainedPolicy]: The policy class corresponding to the given name.
+
+    Raises:
+        NotImplementedError: If the policy with the given name is not implemented.
+    """
+    if name == "pi0":
+        from opentau.policies.pi0.modeling_pi0 import PI0Policy
+
+        return PI0Policy
+    elif name == "pi05":
+        from opentau.policies.pi05.modeling_pi05 import PI05Policy
+
+        return PI05Policy
+    elif name == "value":
+        from opentau.policies.value.modeling_value import ValueFunction
+
+        return ValueFunction
+    else:
+        raise NotImplementedError(f"Policy with name {name} is not implemented.")
+
+
+def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
+    """Creates a policy configuration object based on the policy type.
+
+    Args:
+        policy_type: The type of the policy (e.g., "pi0", "pi05", "value").
+        **kwargs: Keyword arguments to be passed to the configuration class constructor.
+
+    Returns:
+        PreTrainedConfig: An instance of the corresponding policy configuration class.
+
+    Raises:
+        ValueError: If the policy type is not available.
+    """
+    if policy_type == "pi0":
+        return PI0Config(**kwargs)
+    elif policy_type == "pi05":
+        return PI05Config(**kwargs)
+    elif policy_type == "value":
+        return ValueConfig(**kwargs)
+    else:
+        raise ValueError(f"Policy type '{policy_type}' is not available.")
+
+
+def make_policy(
+    cfg: PreTrainedConfig,
+    ds_meta: LeRobotDatasetMetadata | None = None,
+    features: dict[str, FeatureType] | None = None,
+    stats: dict[str, dict[str, np.ndarray]] | None = None,
+    execution_target: Optional[
+        str
+    ] = None,  # None for unified training, "robot" for robot action decoder inference, "cloud" for VLM on cloud inference
+) -> PreTrainedPolicy:
+    """Make an instance of a policy class.
+
+    This function exists because (for now) we need to parse features from either a dataset or an environment
+    in order to properly dimension and instantiate a policy for that dataset or environment.
+
+    Args:
+        cfg: The config of the policy to make. If `pretrained_path` is set, the policy will
+            be loaded with the weights from that path.
+        ds_meta: Dataset metadata to take input/output shapes and statistics to use for
+            (un)normalization of inputs/outputs in the policy. Defaults to None.
+        features: Input and output features. Defaults to None.
+        stats: Dictionary of statistics for normalization. Defaults to None.
+        execution_target: Target for execution. Can be "robot", "cloud", or None.
+            None implies unified training. "robot" implies robot action decoder inference.
+            "cloud" implies VLM on cloud inference. Defaults to None.
+
+    Returns:
+        PreTrainedPolicy: An instance of the created policy.
+
+    Raises:
+        ValueError: If neither or both `ds_meta` and `features` are provided when features are not already set in config.
+        ValueError: If `execution_target` is invalid.
+    """
+    features_already_set = (
+        isinstance(cfg.input_features, dict)
+        and cfg.input_features
+        and isinstance(cfg.output_features, dict)
+        and cfg.output_features
+    )
+    if (bool(ds_meta) + (features is not None) != 1) and not features_already_set:
+        raise ValueError("Exactly one of ds_meta or features must be provided.")
+
+    if execution_target not in ["robot", "cloud", None]:
+        raise ValueError(
+            f"execution_target must be one of ['robot', 'cloud', None], but got {execution_target}."
+        )
+
+    policy_cls = get_policy_class(cfg.type)
+
+    kwargs = {}
+
+    if ds_meta is not None:
+        features = dataset_to_policy_features(ds_meta.features)
+        kwargs["dataset_stats"] = ds_meta.stats
+
+    if not features_already_set:
+        cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
+        cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
+
+    if stats is not None:
+        kwargs["dataset_stats"] = stats
+
+    if execution_target is not None:
+        kwargs["execution_target"] = execution_target
+
+    kwargs["config"] = cfg
+
+    if cfg.pretrained_path:
+        # Load a pretrained policy and override the config if needed (for example, if there are inference-time
+        # hyperparameters that we want to vary).
+        kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+        policy = policy_cls.from_pretrained(**kwargs)
+    else:
+        # Make a fresh policy.
+        policy = policy_cls(**kwargs)
+
+    assert isinstance(policy, nn.Module)
+
+    # policy = torch.compile(policy, mode="reduce-overhead")
+
+    return policy