Upload folder using huggingface_hub

client.py

@@ -28,6 +28,8 @@ class AppEnv(EnvClient[AppAction, AppObservation, AppState]):
             objectsFound=obs_data.get("objectsFound", []),
             reward=obs_data.get("reward", 0.0),
             isDone=obs_data.get("isDone", False),
+            rewardFeedback=obs_data.get("rewardFeedback", []),
+            rewardList=obs_data.get("rewardList", []),
         )
 
         return StepResult(
@@ -48,4 +50,6 @@ class AppEnv(EnvClient[AppAction, AppObservation, AppState]):
             objectsLeft=payload.get("objectsLeft", []),
             objectsFound=payload.get("objectsFound", []),
             ObjectsPresent=payload.get("ObjectsPresent", {}),
+            rewardFeedback=payload.get("rewardFeedback", []),
+            rewardList=payload.get("rewardList", []),
         )

inference.py

@@ -162,7 +162,7 @@ def main() -> None:
         if observation.isDone:
             break
         
-        time.sleep(10000)
+        time.sleep(100)
         
     print(HISTORY)
 

server/app_environment.py

@@ -21,15 +21,27 @@ class AppEnvironment(Environment):
         self._state = self._new_state()
         self._reset_count = 0
 
+    def _coerce_state(self) -> AppState:
+        if isinstance(self._state, AppState):
+            return self._state
+
+        if isinstance(self._state, dict):
+            self._state = AppState(**self._state)
+            return self._state
+
+        self._state = self._new_state()
+        return self._state
+
     def _new_state(self) -> AppState:
         grid, placed = initGrid()
+        grid_shape = (len(grid), len(grid[0]), len(grid[0][0]))
 
         return AppState(
             episode_id=str(uuid4()),
             step_count=0,
             currentGrid=grid,
-            weightedGrid=initWeightedGrid(),
-            objectsLeft=list(OBJECTS.keys()),
+            weightedGrid=initWeightedGrid(grid_shape),
+            objectsLeft=list(placed.keys()),
             objectsFound=[],
             reward=0.0,
             isDone=False,
@@ -53,63 +65,63 @@ class AppEnvironment(Environment):
         )
 
     def step(self, action: AppAction) -> AppObservation:
-        if not isinstance(self._state, AppState):
-            self._state = self._new_state()
+        state = self._coerce_state()
 
-        self._state.step_count += 1
+        if isinstance(action, dict):
+            action = AppAction(**action)
 
+        state.step_count += 1
         reward = 0.0
 
         if action is None:
             reward -= 10.0
             appendRewardFeedback(
-                self._state,
+                state,
                 "No action is of invalid schema or format. Penalty applied.",
                 reward,
             )
             return AppObservation(
-                currentGrid=self._state.currentGrid,
-                positions=self._state.ObjectsPresent,
-                objectsLeft=self._state.objectsLeft,
-                objectsFound=self._state.objectsFound,
-                reward=self._state.reward,
-                isDone=self._state.isDone,
-                rewardFeedback=self._state.rewardFeedback,
-                rewardList=self._state.rewardList,
+                currentGrid=state.currentGrid,
+                positions=state.ObjectsPresent,
+                objectsLeft=state.objectsLeft,
+                objectsFound=state.objectsFound,
+                reward=state.reward,
+                isDone=state.isDone,
+                rewardFeedback=state.rewardFeedback,
+                rewardList=state.rewardList,
             )
 
         if action.isSegmentation and action is not None:
             reward += 10.0
-            appendRewardFeedback(self._state, "Segmentation successful.", reward)
+            appendRewardFeedback(state, "Segmentation successful.", reward)
 
         if action.placement and action is not None:
-            reward += place(action.isSegmentation, action.placement, self._state)
-            appendRewardFeedback(self._state, "Object placed successfully.", reward)
+            reward += place(action.isSegmentation, action.placement, state)
+            appendRewardFeedback(state, "Object placed successfully.", reward)
 
         if action.findObjects and action is not None:
-            reward += findobject(action.isSegmentation, action.findObjects, self._state)
-            appendRewardFeedback(self._state, "Object found successfully.", reward)
+            reward += findobject(action.isSegmentation, action.findObjects, state)
+            appendRewardFeedback(state, "Object found successfully.", reward)
 
-        if len(self._state.objectsLeft) == 0:
-            self._state.isDone = True
+        if len(state.objectsLeft) == 0:
+            state.isDone = True
             reward += 10.0
-            appendRewardFeedback(
-                self._state, "All objects found. Episode completed!", reward
-            )
+            appendRewardFeedback(state, "All objects found. Episode completed!", reward)
 
-        self._state.reward += reward / (10**self._state.step_count)
+        state.reward += reward / (10**state.step_count)
 
         return AppObservation(
-            currentGrid=self._state.currentGrid,
-            positions=self._state.ObjectsPresent,
-            objectsLeft=self._state.objectsLeft,
-            objectsFound=self._state.objectsFound,
-            reward=self._state.reward,
-            isDone=self._state.isDone,
-            rewardFeedback=self._state.rewardFeedback,
-            rewardList=self._state.rewardList,
+            currentGrid=state.currentGrid,
+            positions=state.ObjectsPresent,
+            objectsLeft=state.objectsLeft,
+            objectsFound=state.objectsFound,
+            reward=state.reward,
+            isDone=state.isDone,
+            rewardFeedback=state.rewardFeedback,
+            rewardList=state.rewardList,
         )
 
     @property
     def state(self) -> dict:
-        return self._state.model_dump()
+        state = self._coerce_state()
+        return state.model_dump()

utils.py

@@ -122,8 +122,11 @@ def initGrid():
     return (grid, placed)
 
 
-def initWeightedGrid():
-    grid = random.uniform(0, 1, (randint(5, 11), randint(5, 11), randint(5, 11)))
+def initWeightedGrid(shape=None):
+    if shape is None:
+        shape = (randint(5, 11), randint(5, 11), randint(5, 11))
+
+    grid = random.uniform(0, 1, shape)
 
     x_mid = grid.shape[0] // 2
     x_span = grid.shape[0] // 4
@@ -134,6 +137,23 @@ def initWeightedGrid():
     return grid
 
 
+def _get_weight_value(weight, x, y, z):
+    if not weight or not weight[0] or not weight[0][0]:
+        return 0.0
+
+    if (
+        x < 0
+        or y < 0
+        or z < 0
+        or x >= len(weight)
+        or y >= len(weight[0])
+        or z >= len(weight[0][0])
+    ):
+        return 0.0
+
+    return weight[x][y][z]
+
+
 def place(segment, objects, state):
     dims = state.currentGrid
     weight = state.weightedGrid
@@ -189,17 +209,22 @@ def place(segment, objects, state):
                     elif (
                         dims[pos[0] + i][pos[1] + j][pos[2] + k] > 0 and pos[3] == True
                     ):
-                        if pos[2] + k + 1 <= len(objGrid[0][0]):
+                        if pos[2] + k + 1 < len(dims[0][0]):
                             dims[pos[0] + i][pos[1] + j][pos[2] + k + 1] += 1
-                            reward += (
-                                weight[pos[0] + i][pos[1] + j][pos[2] + k + 1]
+                            bonus = (
+                                _get_weight_value(
+                                    weight,
+                                    pos[0] + i,
+                                    pos[1] + j,
+                                    pos[2] + k + 1,
+                                )
                                 * reward_per_obj_placed
                             )
+                            reward += bonus
                             appendRewardFeedback(
                                 state,
-                                f"Object '{obj_name}' placed with stacking. Bonus: {weight[pos[0] + i][pos[1] + j][pos[2] + k + 1] * reward_per_obj_placed:.2f}",
-                                weight[pos[0] + i][pos[1] + j][pos[2] + k + 1]
-                                * reward_per_obj_placed,
+                                f"Object '{obj_name}' placed with stacking. Bonus: {bonus:.2f}",
+                                bonus,
                             )
                         else:
                             reward -= reward_per_obj_placed
@@ -214,15 +239,14 @@ def place(segment, objects, state):
 
                     else:
                         dims[pos[0] + i][pos[1] + j][pos[2] + k] = 1
-                        reward += (
-                            reward_per_obj_placed
-                            * weight[pos[0] + i][pos[1] + j][pos[2] + k]
+                        bonus = reward_per_obj_placed * _get_weight_value(
+                            weight, pos[0] + i, pos[1] + j, pos[2] + k
                         )
+                        reward += bonus
                         appendRewardFeedback(
                             state,
-                            f"Object '{obj_name}' placed successfully. Bonus: {weight[pos[0] + i][pos[1] + j][pos[2] + k] * reward_per_obj_placed:.2f}",
-                            weight[pos[0] + i][pos[1] + j][pos[2] + k]
-                            * reward_per_obj_placed,
+                            f"Object '{obj_name}' placed successfully. Bonus: {bonus:.2f}",
+                            bonus,
                         )
                 if placement_failed:
                     break