use std::process::{Command,Stdio};
use std::{thread, time, str};
use std::time::{Duration, Instant};
use std::cmp::{max, min};
use std::fs;
use std::fs::File;
use std::path::Path;
use std::string::String;
use std::io::prelude::*;
use std::io::stdout;
use std::net::{TcpStream, SocketAddr};
use std::sync::mpsc::channel;
use std::env;
use std::iter::repeat;
use std::sync::mpsc;
use std::sync::Arc;
use std::f64;
use rustc_serialize::json::{Json, Object};
extern crate glob;
use glob::glob;
extern crate rand;
extern crate x11;
extern crate vnc;
extern crate image;
extern crate rustc_serialize;
use rustc_serialize::json;
extern crate hyper;
use self::hyper::header::{Headers, Authorization, Basic};
extern crate websocket;
use self::websocket::{Message, Sender, Receiver, Client};
use self::websocket::client::request::Url;
use self::websocket::client;
use self::websocket::client::Request;
use self::websocket::header::extensions::Extension;
use std::collections::HashMap;
pub type gym_reward = f64;
pub type gym_done = bool;
pub type gym_range = Vec<usize>;
#[derive(Debug, Clone)]
pub struct GymShape {
pub action_space: gym_range,
pub observation_space: gym_range,
pub reward_max: gym_reward,
pub reward_min: gym_reward
}
#[derive(Debug, Clone)]
pub struct GymState {
pub screen: Vec<u8>
}
pub trait GymMember {
fn start (&mut self, s: &GymShape, t: &GymState) -> ();
fn reward (&mut self, gym_reward, gym_done) -> ();
fn reset (&mut self) -> ();
fn tick (&mut self, &mut Vec<u8>) -> gym_range;
fn close (&mut self) -> ();
}
pub struct GymRemote {
black_screen: bool,
dirty: bool,
id: u32,
fps: u32,
env_id: String,
duration: u64,
record_dst: String,
mode: String,
vnc: Option<vnc::Client>,
rewarder: Option<((self::websocket::client::Sender<websocket::WebSocketStream>, self::websocket::client::Receiver<websocket::WebSocketStream>))>,
state: GymState,
shape: GymShape,
offset_width: usize,
offset_height: usize,
time: Instant,
frame: u32
}
pub struct Gym {
fps: u32,
env_id: String,
max_parallel: u32,
duration: u64,
record_dst: String
}
#[derive(RustcDecodable, RustcEncodable)]
struct RCBody {
seed: u32,
env_id: String,
fps: u32
}
#[derive(RustcDecodable, RustcEncodable)]
struct RCHeaders {
sent_at: u32,
episode_id: u32,
message_id: u32
}
#[derive(RustcDecodable, RustcEncodable)]
struct RewarderCommand {
method: String,
body: RCBody,
headers: RCHeaders
}
const ATARI_HEIGHT: u32 = 262;
const ATARI_WIDTH: u32 = 160;
impl GymRemote {
pub fn recorder_cleanup(&mut self) -> () {
let real_fps = self.frame / max(self.time.elapsed().as_secs() as u32, 1);
Command::new("ffmpeg")
.arg("-r").arg(format!("{}", real_fps))
.arg("-f").arg("image2")
.arg("-i").arg("mov_out/frame_%0d.png")
.arg("-r").arg("24")
.arg("-pix_fmt").arg("yuv420p")
.arg(self.record_dst.clone())
.spawn();
}
pub fn start<T: GymMember>(&mut self, mut agent: T) {
let rx = self.start_rewarder();
self.start_vnc();
let mut agent = self.start_agent(agent);
let mut prev_action = Vec::new();
use x11::keysym::*;
let atari_actions = &[XK_Left, XK_Right, XK_Up, XK_Down, XK_space];
let flash_actions = &[XK_Pointer_Button1, XK_Left, XK_Right, XK_Up, XK_Down, XK_space];
let mut high_score = 0.0;
loop {
if self.time.elapsed().as_secs() > self.duration { break; }
loop {
match rx.try_recv() {
Ok(s) => {
let ms = &s[..];
let ms = Json::from_str(ms).unwrap();
let o = ms.as_object().unwrap();
let ms = o.get("method").unwrap();
let ms = ms.as_string().unwrap();
match ms {
"v0.env.describe" => { }
"v0.env.reward" => {
let headers = o.get("headers").unwrap();
let headers = headers.as_object().unwrap();
let sent_at = headers.get("sent_at").unwrap();
let sent_at = sent_at.as_f64().unwrap();
let message_id = headers.get("message_id").unwrap();
let message_id = message_id.as_u64().unwrap();
let body = o.get("body").unwrap();
let body = body.as_object().unwrap();
let reward = body.get("reward").unwrap();
let reward = reward.as_f64().unwrap();
let done = body.get("done").unwrap();
let done = done.as_boolean().unwrap();
high_score = high_score + reward;
agent.reward( reward, done );
}
"v0.reply.env.reset" => {
println!("finished game with score: {}", high_score);
high_score = 0.0;
agent.reset();
}
_ => { println!("received unknown rewarder message: {}", s); }
}
}
Err(_) => { break; }
}
}
self.sync();
let mut screen_copy = self.state.screen.clone();
let action = agent.tick(&mut screen_copy);
self.render_frame(screen_copy);
let mut vnc = self.vnc.as_mut().unwrap();
if action == prev_action { } else if self.mode=="atari" && action.len()==5 {
for ai in 0..5 {
if prev_action.len()<5 || prev_action[ai]==action[ai] {
} else if action[ai]==0 {
vnc.send_key_event(false, atari_actions[ai]).unwrap();
} else {
vnc.send_key_event(true, atari_actions[ai]).unwrap();
}
}
} else if self.mode=="flash" && action.len()==8 {
let x = action[0];
let y = action[1];
if prev_action.len()==0 || (x!=prev_action[0] && y!=prev_action[1]) {
}
for ai in 2..8 {
if prev_action.len()<8 || prev_action[ai]==action[ai] {
} else if action[ai]==0 {
vnc.send_key_event(false, flash_actions[ai-2]).unwrap();
} else {
vnc.send_key_event(true, flash_actions[ai-2]).unwrap();
}
}
}
prev_action = action.clone();
}
println!("disrupted game with score: {}", high_score);
self.recorder_cleanup();
}
pub fn sync(&mut self) {
self.sync_rewarder();
self.sync_vnc();
}
pub fn start_agent<T: GymMember>(&mut self, mut agent: T) -> T {
agent.start(&self.shape, &self.state);
self.time = Instant::now();
return agent
}
pub fn start_rewarder(&mut self) -> std::sync::mpsc::Receiver<String> {
let ws_url = &format!("ws://127.0.0.1:{}", 15900+self.id)[..];
let url = Url::parse(ws_url).unwrap();
println!("Connecting to rewarder at {}", url);
let mut request = Client::connect(url).unwrap();
let mut auth_header = self::hyper::header::Authorization(
self::hyper::header::Basic {
username: "openai".to_owned(),
password: Some("openai".to_owned())
}
);
request.headers.set(auth_header);
let mut response = request.send().unwrap();
response.validate().unwrap(); let (mut sender, mut receiver) = response.begin().split();
println!("Recorder websocket is now ready to use at {}", 15900+ self.id);
let reset_cmd = RewarderCommand{
method: "v0.env.reset".to_owned(),
body: RCBody {
seed: 1,
env_id: self.env_id.trim().to_string(),
fps: self.fps
},
headers: RCHeaders {
sent_at: 0,
episode_id: 0,
message_id: 0
}
};
let reset_msg = json::encode(&reset_cmd).unwrap();
println!("Send reset json message to rewarder: {}", reset_msg);
let rst_msg = Message::text(String::from(reset_msg));
sender.send_message(&rst_msg);
let (tx, rx) = channel();
thread::spawn(move || {
for msg in receiver.incoming_messages() {
let msg: Result<websocket::message::Message,_> = msg;
match msg {
Ok(message) => {
match message.opcode {
websocket::message::Type::Close => {
println!("Connection to rewarder terminated.");
},
websocket::message::Type::Ping => {
let mut pong_message = Message::pong(message.payload);
sender.send_message(&pong_message);
},
websocket::message::Type::Text => {
let bytes = message.payload.into_owned();
let msg = String::from_utf8(bytes).unwrap();
tx.send(msg).unwrap();
},
_ => {}
}
}
Err(e) => {
}
}
}
});
return rx;
}
pub fn start_vnc(&mut self) {
let vnc_addr: SocketAddr = format!("127.0.0.1:{}", 5900+self.id).parse().expect("Unable to parse socket address");
let stream = match std::net::TcpStream::connect(vnc_addr) {
Ok(stream) => stream,
Err(error) => {
panic!("cannot connect to localhost:{}: {}", 5900+self.id, error);
std::process::exit(1)
}
};
self.vnc = Some(match vnc::Client::from_tcp_stream(stream, true, |methods| {
for method in methods {
match method {
&vnc::client::AuthMethod::Password => {
let mut key = [0; 8];
for (i, byte) in "openai".bytes().enumerate() {
if i == 8 { break }
key[i] = byte
}
return Some(vnc::client::AuthChoice::Password(key))
}
_ => ()
}
}
None
}) {
Ok(vnc) => vnc,
Err(error) => {
panic!("cannot initialize VNC session: {}", error);
std::process::exit(1)
}
});
println!("Connected to vnc on port: {}", 5900+self.id);
}
pub fn sync_rewarder(&mut self) -> () {
}
pub fn render_frame(&mut self, screen: Vec<u8>) {
let width = self.shape.observation_space[0];
let height = self.shape.observation_space[1];
if self.dirty {
let mut imgbuf = image::ImageBuffer::new( width as u32, height as u32 );
for x in 0 .. width {
for y in 0 .. height {
let left = 3*(y * width + x) as usize;
imgbuf.put_pixel(x as u32, y as u32, image::Rgb([ screen[left], screen[left+1], screen[left+2] ]));
}
}
let ref mut fout = File::create(&Path::new( &format!("mov_out/frame_{}.png", self.frame)[..] )).unwrap();
let _ = image::ImageRgb8(imgbuf).save(fout, image::PNG);
self.dirty = false;
} else {
let from_p = format!("mov_out/frame_{}.png", self.frame-1);
let to_p = format!("mov_out/frame_{}.png", self.frame);
std::fs::copy(Path::new(&from_p), Path::new(&to_p));
}
self.frame = self.frame + 1;
}
pub fn sync_vnc(&mut self) -> () {
let width = self.shape.observation_space[0];
let height = self.shape.observation_space[1];
let pause = 1000 / self.fps;
std::thread::sleep_ms(pause);
let mut vnc = self.vnc.as_mut().unwrap();
vnc.request_update(vnc::Rect {
left: self.offset_width as u16,
top: self.offset_height as u16,
width: (self.offset_width+width) as u16,
height: (self.offset_height+height) as u16
}, false).unwrap();
for event in vnc.poll_iter() {
use vnc::client::Event;
match event {
Event::EndOfFrame => {
if !self.black_screen {
self.dirty = true;
}
},
Event::PutPixels(vnc_rect, ref pixels) => {
self.black_screen = false;
for x in vnc_rect.left .. (vnc_rect.left+vnc_rect.width) {
for y in vnc_rect.top .. (vnc_rect.top+vnc_rect.height) {
let i = x - vnc_rect.left;
let j = y - vnc_rect.top;
let left = 4*(j * vnc_rect.width + i) as usize;
if pixels[left]>20 { self.black_screen = false }
if pixels[left+1]>20 { self.black_screen = false }
if pixels[left+2]>20 { self.black_screen = false }
}
}
if !self.black_screen {
for x in vnc_rect.left .. min((self.offset_width+width) as u16, (vnc_rect.left+vnc_rect.width)) {
for y in vnc_rect.top .. min((self.offset_height+height) as u16, (vnc_rect.top+vnc_rect.height)) {
let i = x - vnc_rect.left;
let j = y - vnc_rect.top;
let left = 4*(j * vnc_rect.width + i) as usize;
if (x as usize)>=self.offset_width && (y as usize)>=self.offset_height {
let x = (x as usize) - self.offset_width;
let y = (y as usize) - self.offset_height;
self.state.screen[3*(y*width + x)] = pixels[left+2];
self.state.screen[3*(y*width + x)+1] = pixels[left+1];
self.state.screen[3*(y*width + x)+2] = pixels[left];
}
}
}}
},
_ => {}
}
}
}
}
impl Gym {
pub fn parse_args(&mut self) -> () {
let mut prev = "".to_string();
for argument in std::env::args() {
if prev == "--fps" {
self.fps = argument.parse::<u32>().unwrap();
prev = "".to_string();
} else if prev == "--game" {
self.env_id = argument;
prev = "".to_string();
} else if prev == "--parallel" {
self.max_parallel = argument.parse::<u32>().unwrap();
prev = "".to_string();
} else if prev == "--duration" {
self.duration = argument.parse::<u64>().unwrap();
prev = "".to_string();
} else if prev == "--record" {
self.record_dst = argument;
prev = "".to_string();
} else {
prev = argument;
}
}
}
pub fn set_fps(&mut self, fps: u32) -> () { self.fps = fps }
pub fn set_game(&mut self, game: String) -> () { self.env_id = game.to_string() }
pub fn set_max_parallel(&mut self, par: u32) -> () { self.max_parallel = par }
pub fn set_duration(&mut self, num: u64) -> () { self.duration = num }
pub fn set_record(&mut self, dst: String) -> () { self.record_dst = dst.to_string() }
pub fn new() -> Gym {
Gym {
fps: 10,
env_id: "gym-core.AirRaid-v0".to_string(),
max_parallel: 1,
duration: 180,
record_dst: "video.mpg".to_string()
}
}
pub fn remote_prep_container(&mut self, pi: u32) -> () {
let base_vnc = 5900 + pi;
let base_rec = 15900 + pi;
if self.env_id.starts_with("flashgames.") {
Command::new("docker")
.arg("run")
.arg("-p").arg( format!("5900:{}", base_vnc) )
.arg("-p").arg( format!("15900:{}", base_rec) )
.arg("--ipc")
.arg("host")
.arg("--cap-add")
.arg("SYS_ADMIN")
.arg("--privileged")
.arg("quay.io/openai/universe.flashgames:0.20.25")
.spawn();
} else {
Command::new("docker")
.arg("run")
.arg("-p").arg( format!("5900:{}", base_vnc) )
.arg("-p").arg( format!("15900:{}", base_rec) )
.arg("quay.io/openai/universe.gym-core:0.20.6")
.spawn();
}
print!("spawned docker process at {} / {}", base_vnc, base_rec);
for pi in 0..self.max_parallel {
let mut ok = false;
for _ in 0..15 { if !ok {
let rec_port = 15900 + pi;
println!("Polling docker at port {} for connectivity.", rec_port);
let one_second = time::Duration::from_millis(1000);
thread::sleep(one_second);
let ns1: SocketAddr = format!("127.0.0.1:{}", rec_port).parse().expect("Unable to parse socket address");
let mut s1 = TcpStream::connect(ns1);
match s1 {
Ok(_) => { ok=true; }
_ => { println!("No connectivity to {}", rec_port); }
}
}
}
if !ok { panic!("Unable to confirm connectivity to docker #{}", pi) }
else { println!("Confirmed connectivity to docker #{}", pi); }
}
let ten_seconds = time::Duration::from_millis(10000);
thread::sleep(ten_seconds);
}
pub fn remote_prep_recorder(&mut self, pi: u32) {
for entry in glob("mov_out/*.png").expect("Failed to read glob pattern") {
match entry {
Ok(path) => {
fs::remove_file(path);
}
Err(e) => {}
}
}
for entry in glob("*.mpg").expect("Failed to read glob pattern") {
match entry {
Ok(path) => {
fs::remove_file(path);
}
Err(e) => {}
}
}
fs::create_dir("mov_out/");
}
pub fn start_remote(&mut self, pi: u32) -> GymRemote {
self.remote_prep_container(pi);
self.remote_prep_recorder(pi);
let mut screen_width = ATARI_WIDTH;
let mut screen_height = ATARI_HEIGHT;
let mut offset_width = 0;
let mut offset_height = 0;
let mut mode = "atari";
if self.env_id.starts_with("flashgames.") {
let data = Json::from_str(include_str!("../flashgames.json")).unwrap();
let obj = data.as_object().unwrap()
.get(&self.env_id[..]).unwrap()
.as_object().unwrap();
screen_width = obj.get("width").unwrap().as_u64().unwrap() as u32;
screen_height = obj.get("height").unwrap().as_u64().unwrap() as u32;
offset_width = 20;
offset_height = 84;
mode = "flash";
}
let mut action_space: std::vec::Vec<usize> = Vec::new();
if mode=="flash" {
action_space.push(screen_width as usize);
action_space.push(screen_height as usize);
action_space.push(2);
action_space.push(2); action_space.push(2); action_space.push(2); action_space.push(2); action_space.push(2); } else {
action_space.push(2); action_space.push(2); action_space.push(2); action_space.push(2); action_space.push(2); }
let r = GymRemote {
black_screen: false,
dirty: true,
id: pi,
fps: self.fps,
env_id: format!("{: <99}", self.env_id).clone(),
duration: self.duration,
record_dst: self.record_dst.clone(),
mode: mode.to_string(),
vnc: None,
rewarder: None,
state: GymState {
screen: vec![0; (screen_width * screen_height * 3) as usize]
},
shape: GymShape {
action_space: action_space,
observation_space: vec![screen_width as usize, screen_height as usize, 3 as usize],
reward_max : f64::INFINITY,
reward_min : f64::NEG_INFINITY
},
offset_width: offset_width,
offset_height: offset_height,
time: Instant::now(),
frame: 0
};
return r;
}
pub fn remote_sanitize(&mut self) -> () {
Command::new("sh").arg("-c").arg("docker kill $(docker ps -q); true").spawn();
}
pub fn start<F,T: GymMember>(&mut self, start_agent: F) -> ()
where F: Fn() -> T + Send + Sync + 'static
{
let start_agent = Arc::new(start_agent);
let mut threads = Vec::new();
self.remote_sanitize();
for pi in 0..self.max_parallel {
let start_agent = start_agent.clone();
let mut remote = self.start_remote(pi);
threads.push(std::thread::spawn(move || {
let agent = start_agent();
remote.start(agent);
}));
}
for t in threads {
t.join();
}
self.remote_sanitize();
}
}