Acktr mode in play tool
@@ -3,7 +3,7 @@ import argparse
import gym
import roboschool
-from lib import model
+from lib import model, kfac
from PIL import Image
import numpy as np
@@ -19,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("-s", "--save", type=int, help="If specified, save every N-th step as an image")
+ parser.add_argument("--acktr", default=False, action='store_true', help="Enable Acktr-specific tweaks")
args = parser.parse_args()
env = gym.make(args.env)
@@ -26,6 +27,8 @@ if __name__ == "__main__":
env = gym.wrappers.Monitor(env, args.record)
net = model.ModelActor(env.observation_space.shape[0], env.action_space.shape[0])
+ if args.acktr:
+ opt = kfac.KFACOptimizer(net)
net.load_state_dict(torch.load(args.model))
obs = env.reset()
