Spaces:
Sleeping
Sleeping
| """ | |
| Gymnasium-compatible RL Environment for ChaosLab. | |
| Wraps the TrainingEnv into a standard gymnasium.Env so it can be used | |
| with any RL library (stable-baselines3, ray/rllib, cleanrl, etc.). | |
| The action space is DISCRETE: the agent picks from a curated list of | |
| bash commands that are relevant to SRE troubleshooting scenarios. | |
| The observation space is a text string (the terminal output), encoded | |
| as a Box of ASCII character codes for compatibility with standard RL. | |
| Usage: | |
| from src.rl_env import ChaosLabEnv | |
| env = ChaosLabEnv(scenario="cascading_db_failure") | |
| obs, info = env.reset() | |
| action = env.action_space.sample() # random command | |
| obs, reward, terminated, truncated, info = env.step(action) | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| import gymnasium as gym | |
| from gymnasium import spaces | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from .environment import TrainingEnv | |
| # ====================================================================== | |
| # DISCRETE ACTION CATALOG | |
| # ====================================================================== | |
| # These are the bash commands an RL agent can choose from. | |
| # They cover every scenario in the project: diagnostics, service mgmt, | |
| # file operations, disk/memory/network troubleshooting, and security. | |
| ACTION_CATALOG: List[str] = [ | |
| # ── Diagnostics & Observation ────────────────────────────────── | |
| "ps aux", | |
| "ps", | |
| "top", | |
| "df -h", | |
| "du -sh /var/log/*", | |
| "du -sh /tmp/*", | |
| "free -m", | |
| "uptime", | |
| "whoami", | |
| "hostname", | |
| "date", | |
| "pwd", | |
| # ── Log Investigation ────────────────────────────────────────── | |
| "cat /var/log/app.log", | |
| "cat /var/log/auth.log", | |
| "cat /var/log/nginx/error.log", | |
| "cat /var/log/nginx/access.log", | |
| "cat /var/log/cron.log", | |
| "cat /var/log/syslog", | |
| "tail -n 50 /var/log/app.log", | |
| "tail -n 50 /var/log/auth.log", | |
| "tail -n 50 /var/log/nginx/error.log", | |
| "grep \"500\" /var/log/app.log", | |
| "grep \"error\" /var/log/app.log", | |
| "grep \"Failed\" /var/log/auth.log", | |
| "grep \"Failed password\" /var/log/auth.log", | |
| "grep -i \"error\" /var/log/nginx/error.log", | |
| "grep \"10.99.99.5\" /var/log/auth.log", | |
| "grep \"Accepted\" /var/log/auth.log", | |
| "grep -v \"Failed\" /var/log/auth.log", | |
| # ── Service Management ───────────────────────────────────────── | |
| "systemctl status postgres", | |
| "systemctl restart postgres", | |
| "systemctl start postgres", | |
| "systemctl stop postgres", | |
| "systemctl status nginx", | |
| "systemctl restart nginx", | |
| "systemctl start nginx", | |
| "systemctl status app", | |
| "systemctl restart app", | |
| "systemctl start app", | |
| "systemctl status sshd", | |
| "systemctl restart sshd", | |
| # ── File & Permission Operations ─────────────────────────────── | |
| "ls -la /home/user/scripts/", | |
| "ls -la /home/user/scripts/cleanup.sh", | |
| "ls -la /etc/nginx/sites-enabled/default", | |
| "ls -la /usr/local/bin/backup.sh", | |
| "chmod 0755 /home/user/scripts/cleanup.sh", | |
| "chmod +x /home/user/scripts/cleanup.sh", | |
| "chmod 0755 /usr/local/bin/backup.sh", | |
| "chmod +x /usr/local/bin/backup.sh", | |
| "cat /home/user/scripts/cleanup.sh", | |
| "cat /etc/nginx/sites-enabled/default", | |
| "touch /usr/local/bin/backup.sh", | |
| 'echo "#!/bin/bash" > /usr/local/bin/backup.sh', | |
| # ── Config Repair ────────────────────────────────────────────── | |
| 'echo "server {\n listen 80 default_server;\n server_name _;\n\n location / {\n proxy_pass http://127.0.0.1:8080;\n proxy_set_header Host \\$host;\n proxy_set_header X-Real-IP \\$remote_addr;\n }\n}" > /etc/nginx/sites-enabled/default', | |
| # ── Disk Cleanup ─────────────────────────────────────────────── | |
| "rm /var/log/nginx/error.log", | |
| "rm /var/log/app.log", | |
| "rm /tmp/old_debug.log", | |
| "find /var/log -name '*.log' -type f", | |
| "echo '' > /var/log/nginx/error.log", | |
| "echo '' > /var/log/app.log", | |
| # ── Network Troubleshooting ──────────────────────────────────── | |
| "netstat -tlnp", | |
| "ss -tlnp", | |
| "curl localhost:8080", | |
| "curl localhost:80", | |
| "curl http://127.0.0.1:8080", | |
| "iptables -L", | |
| # ── Cron ─────────────────────────────────────────────────────── | |
| "crontab -l", | |
| "journalctl -u cron", | |
| # ── Memory ───────────────────────────────────────────────────── | |
| "cat /proc/meminfo", | |
| # ── Environment ──────────────────────────────────────────────── | |
| "env", | |
| "history", | |
| # ── Navigation ───────────────────────────────────────────────── | |
| "ls /", | |
| "ls /var/log", | |
| "ls /etc", | |
| "ls /tmp", | |
| "cd /var/log", | |
| "cd /home/user", | |
| "cd /", | |
| ] | |
| # Total number of actions | |
| NUM_ACTIONS = len(ACTION_CATALOG) | |
| # Maximum observation length in characters | |
| MAX_OBS_LENGTH = 4096 | |
| # ====================================================================== | |
| # GYMNASIUM ENVIRONMENT | |
| # ====================================================================== | |
| class ChaosLabEnv(gym.Env): | |
| """ | |
| Gymnasium RL Environment for ChaosLab SRE scenarios. | |
| Action Space: Discrete(N) — index into ACTION_CATALOG | |
| Observation: Box(0, 127, shape=(MAX_OBS_LENGTH,)) — ASCII-encoded terminal output | |
| Reward: Float from TrainingEnv (penalty per step + bonus for progress) | |
| Args: | |
| scenario: Scenario key (e.g. 'cascading_db_failure', 'disk_space_crisis') | |
| max_obs_len: Maximum observation string length (truncated/padded to this) | |
| """ | |
| metadata = {"render_modes": ["human", "ansi"]} | |
| def __init__( | |
| self, | |
| scenario: str = "log_analysis", | |
| max_obs_len: int = MAX_OBS_LENGTH, | |
| render_mode: Optional[str] = None, | |
| ): | |
| super().__init__() | |
| self.scenario = scenario | |
| self.max_obs_len = max_obs_len | |
| self.render_mode = render_mode | |
| # Gymnasium spaces | |
| self.action_space = spaces.Discrete(NUM_ACTIONS) | |
| self.observation_space = spaces.Box( | |
| low=0, high=127, | |
| shape=(self.max_obs_len,), | |
| dtype=np.int32, | |
| ) | |
| # Internal state | |
| self._env: Optional[TrainingEnv] = None | |
| self._last_output: str = "" | |
| self._episode_reward: float = 0.0 | |
| self._command_history: List[str] = [] | |
| # ── Core Gym API ─────────────────────────────────────────────── | |
| def reset( | |
| self, | |
| seed: Optional[int] = None, | |
| options: Optional[Dict[str, Any]] = None, | |
| ) -> Tuple[np.ndarray, Dict[str, Any]]: | |
| """Reset the environment and return initial observation.""" | |
| super().reset(seed=seed) | |
| # Allow overriding scenario via options | |
| scenario = self.scenario | |
| if options: | |
| scenario = options.get("scenario", scenario) | |
| self._env = TrainingEnv(scenario=scenario) | |
| reset_result = self._env.reset() | |
| self._episode_reward = 0.0 | |
| self._command_history = [] | |
| # Build initial observation string | |
| obs_text = self._format_observation(reset_result) | |
| self._last_output = obs_text | |
| info = { | |
| "task_name": reset_result["info"]["task_name"], | |
| "instructions": reset_result["info"]["instructions"], | |
| "max_steps": reset_result["info"]["max_steps"], | |
| "scenario_key": reset_result["info"].get("scenario_key", ""), | |
| "action_catalog_size": NUM_ACTIONS, | |
| } | |
| return self._encode_obs(obs_text), info | |
| def step( | |
| self, action: int | |
| ) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: | |
| """Execute a discrete action (index into ACTION_CATALOG).""" | |
| assert self._env is not None, "Call reset() before step()" | |
| assert 0 <= action < NUM_ACTIONS, f"Invalid action: {action}" | |
| command = ACTION_CATALOG[action] | |
| self._command_history.append(command) | |
| result = self._env.step(command) | |
| reward = result["reward"] | |
| self._episode_reward += reward | |
| terminated = result["done"] and self._env.score >= 1.0 | |
| truncated = result["done"] and not terminated | |
| # Build observation | |
| obs_text = ( | |
| f"$ {command}\n" | |
| f"{result['info'].get('command_output', '')}\n" | |
| f"---\n" | |
| f"Score: {result['info']['task_score']:.2f} | " | |
| f"Step: {result['info']['step']}/{result['info']['max_steps']}" | |
| ) | |
| self._last_output = obs_text | |
| info = { | |
| "command": command, | |
| "command_output": result["info"].get("command_output", ""), | |
| "task_score": result["info"]["task_score"], | |
| "exit_code": result["info"]["exit_code"], | |
| "step": result["info"]["step"], | |
| "max_steps": result["info"]["max_steps"], | |
| "episode_reward": self._episode_reward, | |
| "task_metadata": result["info"].get("task_metadata", {}), | |
| } | |
| if self.render_mode == "human": | |
| self.render() | |
| return self._encode_obs(obs_text), reward, terminated, truncated, info | |
| def render(self) -> Optional[str]: | |
| """Render the last command output.""" | |
| if self.render_mode == "ansi": | |
| return self._last_output | |
| elif self.render_mode == "human": | |
| print(self._last_output) | |
| return None | |
| def close(self): | |
| """Clean up.""" | |
| self._env = None | |
| def hijack_env(self, active_env: TrainingEnv): | |
| """Hijack an actively running TrainingEnv (from websockets or arena)""" | |
| self._env = active_env | |
| self._episode_reward = 0.0 | |
| def get_current_obs(self) -> np.ndarray: | |
| """Grabs the current state for predicting actions on a hijacked env""" | |
| if not self._env: | |
| return np.zeros(self.max_obs_len, dtype=np.int32) | |
| view = self._env._view() | |
| obs_text = ( | |
| f"=== SCENARIO: {self._env.task.nm} ===\n" | |
| f"Instructions:\n{self._env.task.guide()}\n\n" | |
| f"Current Directory: {view['current_directory']}\n" | |
| f"Processes:\n{view['processes']}\n" | |
| f"Files:\n{view['filesystem']}\n" | |
| f"Max Steps: {self._env.limit}\n" | |
| ) | |
| return self._encode_obs(obs_text) | |
| # ── Helpers ──────────────────────────────────────────────────── | |
| def _encode_obs(self, text: str) -> np.ndarray: | |
| """Encode a text string as a fixed-length array of ASCII codes.""" | |
| # Truncate or pad to max_obs_len | |
| text = text[:self.max_obs_len] | |
| codes = [ord(c) if ord(c) < 128 else 32 for c in text] | |
| # Pad with zeros | |
| codes += [0] * (self.max_obs_len - len(codes)) | |
| return np.array(codes, dtype=np.int32) | |
| def _format_observation(reset_result: Dict) -> str: | |
| """Format the initial reset observation as a readable string.""" | |
| obs = reset_result["observation"] | |
| info = reset_result["info"] | |
| return ( | |
| f"=== SCENARIO: {info['task_name']} ===\n" | |
| f"Instructions:\n{info['instructions']}\n\n" | |
| f"Current Directory: {obs['current_directory']}\n" | |
| f"Processes:\n{obs['processes']}\n" | |
| f"Files:\n{obs['filesystem']}\n" | |
| f"Max Steps: {info['max_steps']}\n" | |
| ) | |
| # ── Utility Methods ──────────────────────────────────────────── | |
| def get_action_catalog() -> List[str]: | |
| """Return the full list of available commands.""" | |
| return ACTION_CATALOG.copy() | |
| def action_to_command(action: int) -> str: | |
| """Convert a discrete action index to the corresponding bash command.""" | |
| return ACTION_CATALOG[action] | |
| def command_to_action(command: str) -> Optional[int]: | |
| """Find the action index for a given command string, or None.""" | |
| try: | |
| return ACTION_CATALOG.index(command) | |
| except ValueError: | |
| return None | |
| def get_episode_summary(self) -> Dict[str, Any]: | |
| """Get a summary of the current episode.""" | |
| if self._env is None: | |
| return {} | |
| return { | |
| "task_name": self._env.task.nm, | |
| "difficulty": self._env.difficulty, | |
| "final_score": self._env.score, | |
| "steps_used": self._env.step_count, | |
| "max_steps": self._env.limit, | |
| "completed": self._env.finished, | |
| "episode_reward": self._episode_reward, | |
| "commands_executed": self._command_history.copy(), | |
| "efficiency": round(1 - (self._env.step_count / self._env.limit), 2) | |
| if self._env.finished else 0.0, | |
| } | |
| # ====================================================================== | |
| # REGISTER WITH GYMNASIUM | |
| # ====================================================================== | |
| # Register so users can do: gym.make("ChaosLab-v0", scenario="...") | |
| gym.register( | |
| id="ChaosLab-v0", | |
| entry_point="src.rl_env:ChaosLabEnv", | |
| kwargs={"scenario": "log_analysis"}, | |
| ) | |