mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-26 08:17:18 +00:00
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
"""
|
|
This script does a few rollouts with an environment and writes the data to an npz file
|
|
Its purpose is to help with verifying that you haven't functionally changed an environment.
|
|
(If you have, you should bump the version number.)
|
|
"""
|
|
import argparse, numpy as np, collections, sys
|
|
from os import path
|
|
|
|
|
|
class RandomAgent(object):
|
|
def __init__(self, ac_space):
|
|
self.ac_space = ac_space
|
|
def act(self, _):
|
|
return self.ac_space.sample()
|
|
|
|
def rollout(env, agent, timestep_limit):
|
|
"""
|
|
Simulate the env and agent for timestep_limit steps
|
|
"""
|
|
ob = env.reset()
|
|
data = collections.defaultdict(list)
|
|
for _ in xrange(timestep_limit):
|
|
data["observation"].append(ob)
|
|
action = agent.act(ob)
|
|
data["action"].append(action)
|
|
ob,rew,done,_ = env.step(action)
|
|
data["reward"].append(rew)
|
|
if done:
|
|
break
|
|
return data
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("envid")
|
|
parser.add_argument("outfile")
|
|
parser.add_argument("--gymdir")
|
|
|
|
args = parser.parse_args()
|
|
if args.gymdir:
|
|
sys.path.insert(0, args.gymdir)
|
|
import gym
|
|
from gym import utils
|
|
print utils.colorize("gym directory: %s"%path.dirname(gym.__file__), "yellow")
|
|
env = gym.make(args.envid)
|
|
agent = RandomAgent(env.action_space)
|
|
alldata = {}
|
|
for i in xrange(2):
|
|
np.random.seed(i)
|
|
data = rollout(env, agent, env.spec.timestep_limit)
|
|
for (k, v) in data.items():
|
|
alldata["%i-%s"%(i, k)] = v
|
|
np.savez(args.outfile, **alldata)
|
|
|
|
if __name__ == "__main__":
|
|
main() |