RLHF dataset generation

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Andres Garcia
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -65,6 +65,7 @@ Explore synthetic dataset generation using formal grammars:
 - **`custom_grammar_gen.py`** - Generate datasets using custom grammar configurations
 - **`filtered_dataset_generation.py`** - Advanced dataset generation with simulation-based filtering
 - **`mixed_gen.py`** - Mixed dataset generation combining multiple approaches
+- **`rlhf_dataset_gen.py`** - Interactive RLHF-inspired dataset generation with human-in-the-loop feedback using a unified GUI
 
 #### **Experimental Setup Examples** (`examples/04_experimental_setup_example/`)
 Complete experimental pipeline example for research and evaluation:

agent_control/simulation/agents/agent_sensors.py

@@ -25,7 +25,7 @@ def sense_light(self) -> Vector2:
         Returns:
             Vector2: Direction vector to the nearest light source
         """
-        light_pos = self.agent.env.light_pos
+        light_pos = self.agent.env.light_pos  # type: ignore[attr-defined]
         agent_pos = self.agent.pos
         diff_vec: Vector2 = light_pos - agent_pos
 

agent_control/simulation/agents/elements.py

@@ -11,7 +11,7 @@
     from ..envs.robot_env import RobotEnvironment
 
 
-class Part(Agent):  # type: ignore
+class Part(Agent):
     """Part object that can be picked up and moved by robot agents."""
 
     def __init__(
@@ -38,7 +38,7 @@ def __init__(
         self.simulation = simulation
         self.type = type
         self.pos = pos
-        self.owner: Optional[Agent] = None
+        self.owner: Optional[Agent | "RobotEnvironment"] = None
         self.env = env
         self.is_permanently_placed = False
         self.update_img()
@@ -50,11 +50,11 @@ def update_img(self) -> None:
 
     def remove_part(self) -> None:
         """Remove this part from the simulation."""
-        self.kill()
+        self.kill()  # type: ignore[no-untyped-call]
 
     def update(self) -> None:
         """Update the part's position and state."""
-        if self.owner and not self.is_permanently_placed:
+        if self.owner and not self.is_permanently_placed and hasattr(self.owner, "pos"):
             self.pos = self.owner.pos
 
     def can_be_picked_up(self) -> bool:

agent_control/simulation/agents/robot_agent.py

@@ -5,7 +5,7 @@
 
 import pygame as pg
 from pygame.math import Vector2
-from vi import Agent, Simulation
+from vi import Agent, HeadlessSimulation, Simulation
 
 from tree_parser.middle_parser import parse_behavior_tree
 
@@ -16,13 +16,13 @@
     from ..envs.base_env import SimEnvironment
 
 
-class RobotAgent(Agent):  # type: ignore
+class RobotAgent(Agent):
     """Robot agent that executes behavior trees in simulation environments."""
 
     def __init__(
         self,
         images: List[pg.Surface],
-        simulation: Simulation,
+        simulation: Simulation | HeadlessSimulation,
         pos: Vector2,
         env: "SimEnvironment",
         xml_path: str,

agent_control/simulation/envs/robot_env.py

@@ -88,32 +88,32 @@ def draw_arena(self) -> None:
         """Draw arena boundaries in the simulation environment."""
         self.simulation.spawn_obstacle(
             "./agent_control/simulation/images/arena_new.png",
-            self.arena_pos.x,
-            self.arena_pos.y,
+            int(self.arena_pos.x),
+            int(self.arena_pos.y),
         )
 
     def draw_source(self) -> None:
         """Draw source area where good parts are located."""
         self.simulation.spawn_site(
             "./agent_control/simulation/images/source_green.png",
-            self.source_pos.x,
-            self.source_pos.y,
+            int(self.source_pos.x),
+            int(self.source_pos.y),
         )
 
     def draw_nest(self) -> None:
         """Draw base/nest area where parts should be delivered."""
         self.simulation.spawn_site(
             "./agent_control/simulation/images/blue_nest.png",
-            self.base_pos.x,
-            self.base_pos.y,
+            int(self.base_pos.x),
+            int(self.base_pos.y),
         )
 
     def draw_waste(self) -> None:
         """Draw waste area where bad parts should be disposed."""
         self.simulation.spawn_site(
             "./agent_control/simulation/images/waste_red.png",
-            self.waste_pos.x,
-            self.waste_pos.y,
+            int(self.waste_pos.x),
+            int(self.waste_pos.y),
         )
 
     def spawn_part(self, type: str, pos: Vector2) -> None:
@@ -141,7 +141,7 @@ def remove_part(self, part: Part) -> None:
         Args:
             part: Part object to remove from simulation
         """
-        part.kill()
+        part.kill()  # type: ignore[no-untyped-call]
 
     def place_parts(self, num_parts: int) -> None:
         """

data_grammar/__init__.py

@@ -1,5 +1,6 @@
 """Data Grammar package for generating behavior trees and datasets."""
 
 from .dataset_generator import DatasetGenerator
+from .rlhf_generation.rlhf_unified import UnifiedRLHFUI
 
-__all__ = ["DatasetGenerator"]
+__all__ = ["DatasetGenerator", "UnifiedRLHFUI"]

data_grammar/dataset_generation/sys_prompt.py

@@ -504,7 +504,7 @@
 </BehaviorTree>
 
 Metrics:
-{"good_parts_picked_up": 1}
+{{"good_parts_picked_up": 1}}
 
 Task: Your task is to find a good part and stop moving.
 
@@ -528,7 +528,7 @@
 </BehaviorTree>
 
 Metrics:
-{"good_parts_picked_up": 1}
+{{"good_parts_picked_up": 1}}
 
 Task: Your task is to find a scrap part and stop moving.
 
@@ -552,7 +552,7 @@
 </BehaviorTree>
 
 Metrics:
-{"bad_parts_picked_up": 1}
+{{"bad_parts_picked_up": 1}}
 
 -Examples for finding and bringing parts to areas:
 
@@ -584,7 +584,7 @@
 </BehaviorTree>
 
 Metrics:
-{"good_parts_picked_up": 1, "parts_dropped_in_construction": [1, 0]}
+{{"good_parts_picked_up": 1, "parts_dropped_in_construction": [1, 0]}}
 
 Task: Your task is to find scrap parts and bring them to the source area.
 
@@ -614,7 +614,7 @@
 </BehaviorTree>
 
 Metrics:
-{"bad_parts_picked_up": 1, "parts_dropped_in_source": [0, 1]}
+{{"bad_parts_picked_up": 1, "parts_dropped_in_source": [0, 1]}}
 
 
 - Complex example on how to handle both good and scrap parts to different areas:
@@ -666,7 +666,7 @@
 </BehaviorTree>
 
 Metrics:
-{"good_parts_picked_up": 1, "bad_parts_picked_up": 1, "parts_dropped_in_storage": [1, 0], "parts_dropped_in_source": [0, 1]}
+{{"good_parts_picked_up": 1, "bad_parts_picked_up": 1, "parts_dropped_in_storage": [1, 0], "parts_dropped_in_source": [0, 1]}}
 
 
 You can use the examples above to help you understand how nodes come toguether to build functional units of a behaviour tree, but you can also take direct inspiration to populate a tree if the strucutre is the same or similar to an example.
@@ -722,7 +722,7 @@
 </BehaviorTree>
 
 Metrics:
-{"bad_parts_picked_up": 1}
+{{"bad_parts_picked_up": 1}}
 
 """
 

data_grammar/dataset_generator.py

@@ -27,7 +27,7 @@
 
 # Add type: ignore for dotenv import since it's an optional dependency
 try:
-    from dotenv import load_dotenv  # type: ignore
+    from dotenv import load_dotenv
 
     load_dotenv()
 except ImportError:

data_grammar/rlhf_generation/__init__.py

@@ -0,0 +1 @@
+"""Module for generating datasets trough methods inspired by RLHF."""

data_grammar/rlhf_generation/output/dataset_path.json

@@ -0,0 +1,98 @@
+[
+    {
+        "layman_task": "Find scrap parts and pick them up",
+        "technical_task": "if you detect a scrap part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a scrap part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_scrap_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and take them to the base",
+        "technical_task": "if you are holding a good part then go to the base, or if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you are holding a good part then seek the base area, or if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_agent_holding_good_part</Condition>\n            <StateAction>state_seek_base_area</StateAction>\n        </Sequence>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find scrap parts and take them to the waste",
+        "technical_task": "if you detect a scrap part then pick up the part, or if you are holding a scrap part then go to the waste, or if you are in the waste then drop the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a scrap part then pick up the part, or if you are holding a scrap part then seek the waste area, or if you are in the waste area then drop the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_scrap_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <Sequence>\n            <Condition>is_agent_holding_scrap_part</Condition>\n            <StateAction>state_seek_waste_area</StateAction>\n        </Sequence>\n        <Sequence>\n            <Condition>is_agent_in_waste_area</Condition>\n            <ActuatorAction>drop_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find scrap parts and pick them up",
+        "technical_task": "if you detect a scrap part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a scrap part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_scrap_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up nigga",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Pick up scrap parts and pick them up",
+        "technical_task": "if you detect a scrap part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a scrap part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_scrap_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then go to the source.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then seek the source area.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_seek_source_area</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find good parts and pick them up",
+        "technical_task": "if you detect a good part then pick up the part, or otherwise then search randomly.",
+        "spoon_task": "if you detect a good part then pick up the part, or otherwise then walk randomly.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_good_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n        <StateAction>state_random_walk</StateAction>\n    </Selector>\n</BehaviorTree>"
+    },
+    {
+        "layman_task": "Find scrap parts and picm them up and take them to the waste area and drop them there",
+        "technical_task": "if you are holding a scrap part and if you are in the waste then drop the part, or if you are holding a scrap part then go to the waste, or otherwise if you detect a scrap part then pick up the part.",
+        "spoon_task": "if you are holding a scrap part and if you are in the waste area then drop the part, or if you are holding a scrap part then seek the waste area, or otherwise if you detect a scrap part then pick up the part.",
+        "tree": "<BehaviorTree>\n    <Selector>\n        <Sequence>\n            <Condition>is_agent_holding_scrap_part</Condition>\n            <Condition>is_agent_in_waste_area</Condition>\n            <ActuatorAction>drop_part</ActuatorAction>\n        </Sequence>\n        <Sequence>\n            <Condition>is_agent_holding_scrap_part</Condition>\n            <StateAction>state_seek_waste_area</StateAction>\n        </Sequence>\n        <Sequence>\n            <Condition>is_scrap_part_detected</Condition>\n            <ActuatorAction>pick_up_part</ActuatorAction>\n        </Sequence>\n    </Selector>\n</BehaviorTree>"
+    }
+]