NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-28 17:29:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Create optimizer_benchmark_environmenr.py

Browse source
This commit is contained in:
ParsaIdp 2025-05-18 18:14:10 -07:00 • committed by GitHub
parent f9a444b6f2
commit 71f6d48e87
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
environments/optimizer/optimizer_benchmark_environmenr.py Normal file
View file

@ -0,0 +1,16 @@
from atroposlib.envs.base import BaseEnvConfig, BaseEnv
from atropos.environments.optimizer.wrapper import score_optimizer
class OptimizerBenchmarkEnvConfig(BaseEnvConfig):
architecture: str = "mnist"
class OptimizerBenchmarkEnvironment(BaseEnv):
env_config_cls = OptimizerBenchmarkEnvConfig
def __init__(self, config: OptimizerBenchmarkEnvConfig, server_configs=None, slurm=False, testing=False):
super().__init__(config, server_configs, slurm, testing)
self.architecture = config.architecture
async def evaluate(self, optimizer_code: str, *args, **kwargs):
# This method is required by BaseEnv
return score_optimizer(optimizer_code, self.architecture)