InfraGym

Unified Infrastructure Simulation Environments for LLM Agents

Environments

1. Load Balancing (load_balance)

Optimize job distribution across multiple servers with varying capabilities:

• Goal: Minimize job waiting time and balance server loads
• Actions: Choose which server to assign incoming jobs
• Observations: Server loads, job sizes, service rates, queue states

2. Adaptive Bitrate Streaming (abr_sim)

Manage video streaming quality based on network conditions:

• Goal: Maximize video quality while minimizing stalls and bitrate switches
• Actions: Select video bitrate level (0.3-4.3 Mbps)
• Observations: Buffer size, network throughput, past performance

3. Cache Management (cache)

Decide which objects to cache for optimal hit rates:

• Goal: Maximize cache hit rate while managing storage constraints
• Actions: Accept or reject cache requests
• Observations: Cache state, object metadata, access patterns

Usage Examples

Load Balancing

import infragym

# Create load balancing environment
env = infragym.make('load_balance', 
                   num_servers=3, 
                   max_episode_steps=100,
                   reward_scale=1.0)

obs, info = env.reset()

# Environment provides natural language descriptions
print("Task:", obs['task_description'])
print("Observation:", obs['text_observation'])  
print("Available actions:", obs['available_actions'])

# Take action (assign job to server 1)
obs, reward, done, truncated, info = env.step(1)
print(f"Assigned job to server 1, reward: {reward:.3f}")

env.close()

Adaptive Bitrate Streaming

import infragym

# Create ABR environment
env = infragym.make('abr_sim',
                   max_episode_steps=200,
                   dummy_mode=True)  # Use synthetic network traces

obs, info = env.reset()

# Choose bitrate level (0=lowest, 5=highest)
bitrate_action = 2  # Medium bitrate
obs, reward, done, truncated, info = env.step(bitrate_action)

print(f"Selected bitrate level {bitrate_action}")
print(f"Buffer: {info['buffer_size']:.1f}s")
print(f"Stall time: {info['stall_time']:.3f}s")

env.close()

Cache Management

import infragym

# Create cache environment  
env = infragym.make('cache',
                   cache_size=1000,
                   max_episode_steps=500)

obs, info = env.reset()

# Make caching decision (1=cache, 0=reject)
cache_decision = 1  # Accept the object
obs, reward, done, truncated, info = env.step(cache_decision)

print(f"Cache decision: {'Accept' if cache_decision else 'Reject'}")
print(f"Cache hit rate: {info.get('hit_rate', 0):.2%}")

env.close()

API Reference

Core Functions

import infragym

# Create environment
env = infragym.make(env_name, **kwargs)

# Environment discovery
environments = infragym.list_environments()
info = infragym.get_env_info(env_name)

# Resource cleanup
infragym.close_all()

Environment Parameters

Load Balance (load_balance)

• num_servers: Number of servers (default: 5)
• max_episode_steps: Episode length (default: 50)
• reward_scale: Reward scaling factor (default: 1.0)

ABR Streaming (abr_sim)

• max_episode_steps: Episode length (default: 100)
• dummy_mode: Use synthetic traces (default: True)
• reward_scale: Reward scaling factor (default: 1.0)

Cache Management (cache)

• cache_size: Cache capacity (default: 1000)
• max_episode_steps: Episode length (default: 200)
• trace_type: Trace data source (default: 'synthetic')

Observation Format

All environments return observations with:

{
    'text_observation': str,      # Natural language description
    'structured_observation': np.ndarray,  # Numerical features  
    'available_actions': str,     # Action descriptions
    'task_description': str       # Environment-specific instructions
}