Save screenshots
@@ -4,6 +4,7 @@ import gym
import roboschool
from lib import model
+from PIL import Image
import numpy as np
import torch
@@ -18,6 +19,7 @@ if __name__ == "__main__":
parser.add_argument("-e", "--env", default=ENV_ID, help="Environment name to use, default=" + ENV_ID)
parser.add_argument("-r", "--record", help="If specified, sets the recording dir, default=Disabled")
parser.add_argument("-k", "--kind", choices=['a2c', 'd4pg'], default='a2c', help="Kind of model to load. Options: a2c or d4pg")
+ parser.add_argument("-s", "--save", type=int, help="If specified, save every N-th step as an image")
args = parser.parse_args()
env = gym.make(args.env)
@@ -43,4 +45,8 @@ if __name__ == "__main__":
total_steps += 1
if done:
break
+ if args.save is not None and total_steps % args.save == 0:
+ o = env.render('rgb_array')
+ img = Image.fromarray(o)
+ img.save("img_%05d.png" % total_steps)
print("In %d steps we got %.3f reward" % (total_steps, total_reward))
