Initial commit

skillopt/envs/alfworld/vendor/env_base.py

@@ -0,0 +1,84 @@
+# Vendored from SkillRL (Apache-2.0 License)
+# Original: agent_system/environments/base.py
+# Trimmed to only include what ALFWorld needs.
+
+from typing import List, Tuple, Dict, Any
+import numpy as np
+from collections import defaultdict
+
+
+def to_numpy(data):
+    """Convert data to numpy array."""
+    # Lazy-check for torch.Tensor to avoid hard dependency on torch
+    _torch_tensor = None
+    try:
+        import torch
+        _torch_tensor = torch.Tensor
+    except ImportError:
+        pass
+
+    if _torch_tensor is not None and isinstance(data, _torch_tensor):
+        data = data.detach().cpu().numpy()
+    elif isinstance(data, np.ndarray):
+        pass
+    elif isinstance(data, (int, float, bool, Tuple, List)):
+        data = np.array(data)
+    else:
+        raise ValueError(f"Unsupported type: {type(data)})")
+    return data
+
+
+class EnvironmentManagerBase:
+    """Base class for vectorized environment managers.
+
+    Manages a set of parallel environments, handles action projection,
+    observation post-processing, and history tracking.
+    """
+
+    def __init__(self, envs, projection_f, config):
+        self.envs = envs
+        self.projection_f = projection_f
+        self.config = config
+
+    def reset(self, kwargs) -> Dict[str, Any]:
+        obs, infos = self.envs.reset()
+        return {'text': None, 'image': obs, 'anchor': None}, infos
+
+    def step(self, text_actions: List[str]):
+        actions, valids = self.projection_f(text_actions)
+        next_obs, rewards, dones, infos = self.envs.step(actions)
+
+        next_observations = {
+            'text': None,
+            'image': next_obs,
+            'anchor': None,
+        }
+        for i, info in enumerate(infos):
+            info['is_action_valid'] = to_numpy(valids[i])
+
+        rewards = to_numpy(rewards)
+        dones = to_numpy(dones)
+        return next_observations, rewards, dones, infos
+
+    def close(self) -> None:
+        self.envs.close()
+
+    def success_evaluator(self, *args, **kwargs) -> Dict[str, np.ndarray]:
+        total_infos = kwargs['total_infos']
+        total_batch_list = kwargs['total_batch_list']
+        batch_size = len(total_batch_list)
+
+        success = defaultdict(list)
+        for bs in range(batch_size):
+            self._process_batch(bs, total_batch_list, total_infos, success)
+        assert len(success['success_rate']) == batch_size
+        return {key: np.array(value) for key, value in success.items()}
+
+    def _process_batch(self, batch_idx, total_batch_list, total_infos, success):
+        for i in reversed(range(len(total_batch_list[batch_idx]))):
+            batch_item = total_batch_list[batch_idx][i]
+            if batch_item['active_masks']:
+                info = total_infos[batch_idx][i]
+                won_value = float(info['won'])
+                success['success_rate'].append(won_value)
+                return