Team Ai
Apppublic

codecrypt112/openenv-hackathon-ctrlaltwin-tiffenpacker

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
tiffin_environment.py280 linesDownload Raw Back to server
1# Copyright (c) 2026 CtrlAltWin Team2"""3Tiffin Packing Environment — OpenEnv Server Implementation.4 5Wraps the packing simulation into the OpenEnv Environment base class,6exposing step(), reset(), and state() for LLM agent interaction.7"""8 9from __future__ import annotations10 11from typing import Any, Optional12from uuid import uuid413 14try:15    from openenv.core.env_server import Environment16except ImportError:17    # Fallback for local testing without openenv installed18    class Environment:19        def __init__(self, **kwargs): pass20        def reset(self, **kwargs): raise NotImplementedError21        def step(self, action, **kwargs): raise NotImplementedError22        @property23        def state(self): raise NotImplementedError24 25from tiffin_packer.models import TiffinAction, TiffinObservation, TiffinState26from tiffin_packer.simulation.engine import PackingSimulation27from tiffin_packer.vlm.classifier import FoodClassifier28from tiffin_packer.tasks import get_task_config, list_tasks29from tiffin_packer.grader import grade, grade_detailed30 31 32class TiffinPackingEnvironment(Environment):33    """34    OpenEnv-compliant tiffin packing environment.35 36    An LLM agent controls a robotic arm to identify food items using VLM37    and pack them into the correct tiffin containers under real-world38    constraints (type compatibility, volume, temperature, fragility).39 40    Supports 3 tasks: easy, medium, hard.41    """42 43    _instance = None44 45    def __new__(cls, *args, **kwargs):46        if cls._instance is None:47            cls._instance = super().__new__(cls)48            cls._instance._initialized = False49        return cls._instance50 51    def __init__(self):52        if getattr(self, "_initialized", False):53            return54        super().__init__()55        self.sim = PackingSimulation()56        self.vlm = FoodClassifier()57        self._state = TiffinState()58        self._identified_items: set = set()59        self._task_config = None60        self._initialized = True61 62    def reset(63        self,64        seed: Optional[int] = None,65        episode_id: Optional[str] = None,66        **kwargs: Any,67    ) -> TiffinObservation:68        """69        Reset the environment for a new episode.70 71        Args:72            seed: Optional random seed for reproducibility.73            episode_id: Optional custom episode ID.74            **kwargs: Must include 'task_id' (easy/medium/hard).75 76        Returns:77            Initial TiffinObservation with scene description.78        """79        task_id = kwargs.get("task_id", "easy")80 81        # Load task configuration82        self._task_config = get_task_config(task_id, seed=seed)83 84        # Reset simulation85        self.sim.reset(86            food_items=self._task_config.food_items,87            containers=self._task_config.containers,88            seed=seed,89        )90 91        # Reset state92        self._state = TiffinState(93            episode_id=episode_id or str(uuid4()),94            step_count=0,95            task_id=task_id,96            items_packed=0,97            total_items=len(self._task_config.food_items),98            items_identified=0,99            packing_log=[],100            constraints_violated=[],101        )102        self._identified_items = set()103 104        # Build initial observation105        return self._build_observation(106            reward=0.0,107            done=False,108            feedback=(109                f"Episode started! Task: {task_id.upper()}\n\n"110                f"{self._task_config.description}\n\n"111                f"You have {self._task_config.max_steps} steps to pack "112                f"{len(self._task_config.food_items)} food items into "113                f"{len(self._task_config.containers)} containers.\n\n"114                f"Start by using 'observe' to see the scene, then 'identify' "115                f"each food item before packing."116            ),117        )118 119    def step(120        self,121        action: TiffinAction,122        timeout_s: Optional[float] = None,123        **kwargs: Any,124    ) -> TiffinObservation:125        """126        Execute one step in the environment.127 128        Args:129            action: TiffinAction with command and optional target_id.130            timeout_s: Optional timeout (unused).131 132        Returns:133            TiffinObservation with updated scene state.134        """135        self._state.step_count += 1136        reward = 0.0137        done = False138        vlm_result = None139        feedback = ""140 141        command = action.command.lower().strip()142        target_id = action.target_id143 144        # --- Dispatch command ---145        if command == "observe":146            _, feedback, reward = self.sim.observe()147 148        elif command == "identify":149            if target_id is None:150                feedback = "Error: 'identify' requires a target_id (food item ID)."151                reward = -0.1152            else:153                success, feedback, reward, vlm_result = self.sim.identify(target_id)154                if success and vlm_result and vlm_result.get("name"):155                    self._identified_items.add(target_id)156                    self._state.items_identified = len(self._identified_items)157 158        elif command == "pick":159            if target_id is None:160                feedback = "Error: 'pick' requires a target_id (food item ID)."161                reward = -0.1162            else:163                success, feedback, reward = self.sim.pick(target_id)164 165        elif command == "place":166            if target_id is None:167                feedback = "Error: 'place' requires a target_id (container ID)."168                reward = -0.1169            else:170                success, feedback, reward = self.sim.place(target_id)171                if success:172                    self._state.items_packed = sum(173                        1174                        for i in self.sim.food_items175                        if i.status == "packed"176                    )177                    self._state.packing_log = list(self.sim.packing_log)178 179        elif command == "pour":180            if target_id is None:181                feedback = "Error: 'pour' requires a target_id (container ID)."182                reward = -0.1183            else:184                success, feedback, reward = self.sim.pour(target_id)185                if success:186                    self._state.items_packed = sum(187                        1188                        for i in self.sim.food_items189                        if i.status == "packed"190                    )191                    self._state.packing_log = list(self.sim.packing_log)192 193        else:194            feedback = (195                f"Unknown command: '{command}'. "196                f"Available commands: {self.sim.get_available_commands()}"197            )198            reward = -0.1199 200        # --- Time penalty ---201        reward -= 0.02202 203        # --- Check termination ---204        done = (205            self.sim.all_packed206            or self._state.step_count >= self._task_config.max_steps207        )208 209        # --- Final grading ---210        final_score = None211        grade_breakdown = None212        if done:213            grade_breakdown = grade_detailed(214                self._state.packing_log, self._task_config215            )216            final_score = grade_breakdown["final_score"]217            reward += final_score  # bonus = final grade218 219            if self.sim.all_packed:220                feedback += f"\n\n🎉 All items packed! Final score: {final_score:.4f}"221            else:222                feedback += (223                    f"\n\n⏰ Time's up! {self.sim.unpacked_count} items remaining. "224                    f"Final score: {final_score:.4f}"225                )226 227        return self._build_observation(228            reward=reward,229            done=done,230            feedback=feedback,231            vlm_result=vlm_result,232            final_score=final_score,233            grade_breakdown=grade_breakdown,234        )235 236    @property237    def state(self) -> TiffinState:238        """Return the current episode state."""239        return self._state240 241    # -------------------------------------------------------------------242    # Helpers243    # -------------------------------------------------------------------244 245    def _build_observation(246        self,247        reward: float = 0.0,248        done: bool = False,249        feedback: str = "",250        vlm_result: dict = None,251        final_score: float = None,252        grade_breakdown: dict = None,253    ) -> TiffinObservation:254        """Build a TiffinObservation from current state."""255        metadata = {}256        if final_score is not None:257            metadata["final_score"] = final_score258        if grade_breakdown is not None:259            metadata["grade_breakdown"] = grade_breakdown260 261        return TiffinObservation(262            done=done,263            reward=round(reward, 4),264            metadata=metadata,265            scene_description=self.sim.get_scene_description(),266            food_items=[267                item.to_dict(hide_unidentified=True)268                for item in self.sim.food_items269            ],270            containers=[c.to_dict() for c in self.sim.containers],271            held_item=(272                self.sim.held_item.to_dict(hide_unidentified=False)273                if self.sim.held_item274                else None275            ),276            vlm_result=vlm_result,277            available_commands=self.sim.get_available_commands(),278            step_feedback=feedback,279        )280