SafeOR-Gym: A Benchmark Suite for Safe Reinforcement Learning Algorithms on Practical Operations Research Problems

SafeOR-Gym is a benchmark suite of Gym-compatible environments for safe reinforcement learning (SafeRL) in industrially relevant operations research (OR) problems. It is designed to evaluate SafeRL algorithms on realistic, structured, and safety-critical decision-making problems commonly encountered in industrial planning and real-time control.

This suite includes nine environments that model some well-known and challenging problems such as unit commitment, plant scheduling, resource allocation, supply chain logistics, and energy system operations. Each environment integrates strict constraints and planning horizons—making them ideal for testing the safety, robustness, and feasibility performance of RL agents. SafeOR-Gym is natively compatible with the OmniSafe framework, providing out-of-the-box support for constraint-handling algorithms, parallel training, and standardized benchmarking.

The key contributions of this project:

• A modular suite of nine OR-inspired SafeRL environments with varying structures, horizons, and complexities.

• Ready-to-use integration with OmniSafe, enabling immediate use of a large number of SafeRL algorithms.

---

Table of Contents

• Installation
• Usage
• Environments
• Benchmarking Setup (ExperimentGrid)
• Cite us
• License
• How to Contribute

---

Installation

Prerequisites

• Python 3.10
• PyTorch ≥ 1.10
• (Optional) Gurobi / CPLEX for optimization

Install

git clone https://github.com/li-group/SafeOR-Gym.git
cd SafeOR-Gym
git checkout v0.1   
conda env create -f environment.yml
conda activate safeorenv  
pip install -e .

Note: This benchmarking part of the repository uses a modified version of OmniSafe, which includes a few minor changes. The core algorithm implementations remain unchanged.

Usage

The package models SafeRL Environments as CMDP classes (similar to gym). Creating environment instances is very simple:

import SafeOR_Gym
env_rtn = SafeOR_Gym.safeor_make('rtn-v0')
#If required a config file can be added
env_stn = SafeOR_Gym.safeor_make('stn-v0',config_file_path) 

The environment can be then be used with other SafeRL libraries that accept CMDP classes. Furthermore importing the package automatically adds the env-id to the Omnisafe registry making it easy to access Omnisafe algorithms.

Environments

• Production Scheduling in Air Separation Unit (ASUEnv): Optimize liquid production to minimize electricity and production costs, while fulfilling demand and respecting unit capacities across time.
• Generation and Transmission Expansion Planning (GTEPEnv): Plan capacity expansion in power systems under long-term investment and operational constraints.
• Grid Integrated Energy Storage (GridStorageEnv): Manage storage dispatch in a grid setting with price arbitrage and safety limits.
• Integrated Scheduling and Maintenance: Jointly optimize production schedules and maintenance windows under equipment availability constraints.
• Multi-Echelon Supply Chain (InvMgmtEnv): Simulate inventory dynamics across multiple tiers of a supply chain network.
• Multiperiod Blending Problem (BlendingEnv): Solve a multi-time-step blending optimization under ratio, availability, and demand constraints.
• Resource Task Network: Schedule resource-consuming tasks across time with bounded inventories and task delays.
• State Task Network: Model discrete-time transitions of material states via tasks executed on shared units.
• Unit Commitment: Optimize on/off decisions for generators over time while meeting demand and respecting ramping and reserve constraints.

Each environment has its own folder containing the relevant code. To run and benchmark an environment, execute the corresponding script located within its folder.

Environment	env_id(s)
Production Scheduling in Air Separation Unit	ASU1
Generation & Transmission Expansion	Capacity-Expansion
Grid Integrated Energy Storage	Battery-v0
Integrated Scheduling and Maintenance	GASU-v0, GASU-v1
Multi-Echelon Supply Chain	SupplyChain-v0
MultiPeriod Blending	Blending-simple
Resource Task Network	rtn-v0
State Task Network	stn-v0
Unit Commitment	UC-v0, UC-v1

Benchmarking

Standard OmniSafe

SafeOR-Gym environments are compatible with standard OmniSafe workflows. You can run a standard benchmark using the Benchmark_main.py:

python Benchmark_main.py --env_id Battery-v0 --episodes_per_epoch 1 --total_epochs 1

Available Arguments

The Benchmark_main.py script supports the following arguments:

Required:

• --env_id (str): Environment ID (e.g., Battery-v0)
• --episodes_per_epoch (int): Number of episodes per epoch
• --total_epochs (int): Total number of training epochs

Optional:

• --seed (int, default=10): Random seed
• --use_wandb: Enable Weights & Biases logging
• --use_tensorboard: Enable TensorBoard logging (default: enabled)
• --vector_env_nums (int, default=1): Number of parallel environments
• --torch_threads (int, default=1): Number of CPU threads
• --device (str): Device (e.g., cpu, cuda:0) (auto-selected if not provided)
• --gpu_id (int, default=0): GPU ID to use (if CUDA is available)
• --num_pool (int, default=1): ExperimentGrid worker pool size
• --compare_num (int, default=5): Number of algorithms compared during analysis
• --num_episodes (int, default=10): Number of episodes for evaluation

Extended version of OmniSafe

To provide additional control over configurations, we include an extended version of OmniSafe that retains the core algorithm implementations unchanged. We also provide a main.py script to run experiments using this setup. This configuration is used for the experiments reported in the paper.

This setup enables additional flexibility in configuring model and training parameters.

Additional Arguments

The main.py script supports the following arguments:

Required:

• --env_id (str): Environment ID (e.g., Battery-v0)
• --steps_per_epoch (int): Number of steps per epoch
• --total_epochs (int): Total number of training epochs

Optional:

• --environment_config_file_path (str): Path to environment configuration file
• --output_activation_function (str): Actor output activation (e.g., tanh)
• --seed (int, default=10): Random seed
• --use_wandb: Enable Weights & Biases logging
• --use_tensorboard: Enable TensorBoard logging (default: enabled)
• --vector_env_nums (int, default=1): Number of parallel environments
• --torch_threads (int, default=1): Number of CPU threads
• --device (str): Device (e.g., cpu, cuda:0) (auto-selected if not provided)
• --gpu_id (int, default=0): GPU ID to use
• --num_pool (int, default=1): ExperimentGrid worker pool size
• --compare_num (int, default=5): Number of algorithms compared
• --num_episodes (int, default=10): Number of episodes for evaluation

Cite us

Cite us ❤️

@article{ramanujam2025safeor,
  title={SafeOR-Gym: A Benchmark Suite for Safe Reinforcement Learning Algorithms on Practical Operations Research Problems},
  author={Ramanujam, Asha and Elyoumi, Adam and Chen, Hao and Kompalli, Sai Madhukiran and Ahluwalia, Akshdeep Singh and Pal, Shraman and Papageorgiou, Dimitri J and Li, Can},
  journal={arXiv preprint arXiv:2506.02255},
  year={2025}
}

License

This repository is licensed under the MIT License.

How to Contribute

We welcome contributions and bug reports. If you'd like to extend the environments or benchmark new algorithms for any environment, please open a pull request or issue.

Creating New Environments

SafeOR-Gym follows a modular design pattern that makes it easy to create new environments. Below is a step-by-step guide using the Multi-Echelon Supply Chain environment as an example.

1. Base Gym Environment Structure

Create a standard Gymnasium environment by inheriting from gym.Env. Here's the basic structure:

import gymnasium as gym
import numpy as np
from gymnasium.spaces import Box, Dict

class InvMgmtEnv(gym.Env):
    def __init__(self, env_id: str = 'InvMgmt-v0', **kwargs):
        super().__init__()
        
        # Load configuration
        config_path = kwargs.pop('config_path', None)
        raw_cfg = self.load_config(config_path)
        assign_env_config(self, raw_cfg)
        
        # Define observation and action spaces
        self.observation_space = Box(low=low_obs, high=high_obs, shape=(obs_dim,))
        self.action_space = Box(low=-1.0, high=1.0, shape=(act_dim,))
        
    def reset(self, seed=None, options=None):
        # Initialize environment state
        return self._get_state(), {}
        
    def step(self, action):
        # Execute one time step
        # Return: observation, reward, terminated, truncated, info
        pass

2. State Space Definition

Define your state representation with clear structure. For the supply chain example:

def _get_state(self, mode='arr'):
    """Return current state as dict or flattened array"""
    state_dict = {
        'on_hand_inventory': {node: inventory_level},
        'pipeline_inventory': {(i,j): [transit_quantities]},
        'sales': {(retailer, market): units_sold},
        'backlog': {(retailer, market): unfulfilled_demand},
        'demand_window': {(retailer, market): [future_demands]},
        't': current_time_period
    }
    
    if mode == 'dict':
        return state_dict
    else:
        # Flatten for neural network input
        flat_obs, mapping = flatten_and_track_mappings(state_dict)
        return flat_obs

3. Action Space Design

Design continuous or discrete action spaces based on your problem:

# Continuous actions scaled from [-1,1] to actual ranges
def decode_action(self, raw_action):
    action_dict = {}
    for i, route in enumerate(self.reordering_routes):
        # Scale from [-1,1] to [0, capacity]
        scaled_value = (raw_action[i] + 1.0) * 0.5 * self.route_capacity[route]
        action_dict[route] = max(0.0, scaled_value)  # Ensure non-negative
    return action_dict

4. Constraint Handling

Implement constraint violations as costs for safe RL:

def check_action_bounds_cost(self, action_dict):
    """Check action constraints and calculate penalties"""
    penalty = 0.0
    for route, value in action_dict.items():
        # Lower bound constraint
        if value < 0.0:
            penalty += abs(value) * self.penalty_factors['action']
            action_dict[route] = 0.0
        
        # Upper bound constraint (capacity)
        if value > self.route_capacity[route]:
            excess = value - self.route_capacity[route]
            penalty += excess * self.penalty_factors['action']
            action_dict[route] = self.route_capacity[route]
    
    return action_dict, penalty

def check_obs_bounds_cost(self, observation):
    """Check state constraints and calculate penalties"""
    penalty = 0.0
    for i, value in enumerate(observation):
        category, _ = self.obs_mapping[i]
        if category in self.penalty_factors:
            # Inventory capacity constraints
            if value > self.obs_space.high[i]:
                excess = value - self.obs_space.high[i]
                penalty += excess * self.penalty_factors[category]
    
    return penalty

5. Safe RL Wrapper

Create a CMDP wrapper for integration with OmniSafe:

from omnisafe.envs.core import CMDP, env_register

@env_register
class SupplyChainSafe(CMDP):
    _support_envs = ['SupplyChain-v0']
    need_auto_reset_wrapper = True
    need_time_limit_wrapper = True

    def __init__(self, env_id: str, **kwargs):
        super().__init__(env_id)
        self._env = InvMgmtEnv(env_id=env_id, **kwargs.get('env_init_cfgs', {}))
        self._action_space = self._env.action_space
        self._observation_space = self._env.observation_space

    def step(self, action):
        obs, reward, terminated, truncated, info = self._env.step(action.cpu().numpy())
        cost = self._env.cost  # Constraint violations
        return (torch.tensor(obs), torch.tensor(reward), 
                torch.tensor(cost), torch.tensor(terminated), 
                torch.tensor(truncated), info)

6. Configuration Files

Use JSON configuration files to make environments easily customizable:

{
    "T": 30,
    "num_markets": 1,
    "num_retailers": 1,
    "num_distributors": 2,
    "initial_inv": {"1": 100, "2": 120},
    "inventory_holding_cost": {"1": 0.04, "2": 0.03},
    "reordering_route_capacity": {"(2,1)": 500, "(3,1)": 400},
    "penalty_factors": {
        "action": 10.0,
        "inventory": 5.0,
        "pipeline": 3.0
    }
}

7. Integration with Benchmarking

Add your environment to the benchmarking framework:

# In your main benchmarking script
eg = ExperimentGrid(exp_name='Benchmark_YourEnvironment')
eg.add('env_id', ['YourEnv-v0'])
eg.add('env_cfgs:env_init_config:config_file', ['/path/to/config.json'])

# Run experiments
eg.run(train, num_pool=1)

Key Design Principles

1. Modular Structure: Separate the base environment from the safe RL wrapper
2. Configuration-Driven: Use JSON/YAML files for environment parameters
3. Constraint-Aware: Implement constraints as costs, not hard boundaries
4. Scalable State/Action: Design spaces that can handle different problem sizes
5. Clear Documentation: Document state variables, action meanings, and constraints

Testing Your Environment

1. Test basic functionality: env.reset(), env.step(), action/observation spaces
2. Verify constraint handling: Check that violations produce appropriate costs
3. Run with SafeRL algorithms: Ensure compatibility with OmniSafe
4. Validate against domain knowledge: Verify behavior matches expected OR problem dynamics

For more detailed examples, examine the existing environments in the repository, particularly the Multi-Echelon Supply Chain (supply_chain_gym.py, supply_chain_safe.py) and configuration files.

Customizing Constraints

Modifying Existing Constraints

You can easily modify constraints without changing code by editing configuration files:

{
    "penalty_factors": {
        "action": 10.0,        // Penalty for action bound violations
        "inventory": 5.0,      // Penalty for inventory capacity violations
        "pipeline": 3.0        // Penalty for pipeline constraints
    },
    "reordering_route_capacity": {
        "(2,1)": 500,          // Maximum reorder quantity for route (2,1)
        "(3,1)": 400
    },
    "inv_capacity": {
        "1": 1000,             // Maximum inventory at node 1
        "2": 800
    }
}

Adding New Constraint Types

Extend constraint handling by modifying the base environment:

def check_custom_constraints(self, action, state):
    """Add domain-specific constraints"""
    penalty = 0.0
    
    # Example: Production rate constraints
    for producer in self.producers:
        production_rate = action.get(f'produce_{producer}', 0)
        if production_rate > self.max_production_rate[producer]:
            penalty += (production_rate - self.max_production_rate[producer]) ** 2
    
    # Example: Resource availability constraints
    total_resource_usage = sum(action.values())
    if total_resource_usage > self.available_resources:
        penalty += (total_resource_usage - self.available_resources) * 100
    
    return penalty

Community Contributions

We encourage researchers to contribute:

1. New Environments: Submit OR problems from different domains
2. Algorithm Implementations: Add new SafeRL algorithms to the benchmark
3. Evaluation Metrics: Propose domain-specific performance measures
4. Real-World Validation: Provide datasets or case studies for validation

For contributing guidelines, see our GitHub repository and submit pull requests with detailed documentation.