<file_sep>"""
Tensorflow ops for robotic prior losses.
"""
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import ops
import os
from time import strftime

k = {
    "epsilon_a": 0.01,
    "epsilon_r": 0.01,
    "range_W": 0.05,
    "w_temp": 10,
    "w_prop": 1,
    "w_caus": 1,
    "w_rep": 5,
    "lambda": 0.0025,
    "dim_s": 2,
    "num_pair_samples": 100
}

class RoboticPriors:

    def __init__(self, dim_image):
        self.dim_o = np.prod(dim_image)
        self.dim_x = 2
        self.dim_a = 2
        self.lr = 0.001

        self.w_temp = 5
        self.w_prop = 1
        self.w_caus = 5
        self.w_rep = 5

        self.add_placeholders_op()
        self.add_model_op()
        self.add_loss_op()
        self.add_optimizer_op()
        self.add_summary_op()

        self.train_writer = None

    def add_placeholders_op(self):
        # Placeholders
        self.o = tf.placeholder(tf.float32, [None, self.dim_o], "o")
        self.a = tf.placeholder(tf.float32, [None, self.dim_a], "a")
        self.r = tf.placeholder(tf.float32, [None], "r")
        self.x = tf.placeholder(tf.float32, [None, self.dim_x], "x")

    def add_model_op(self):
        dim_h1 = 16
        dim_h2 = 8
        krn_h1 = 5
        krn_h2 = 5

        # Model architecture
        initializer = tf.contrib.layers.xavier_initializer()
        # W1 = tf.get_variable("W1", [dim_o, dim_h1], dtype=tf.float32, initializer=initializer)
        # W2 = tf.get_variable("W2", [dim_h1, dim_x], dtype=tf.float32, initializer=initializer)
        # b1 = tf.get_variable("b1", [dim_h1], dtype=tf.float32, initializer=initializer)
        # b2 = tf.get_variable("b2", [dim_x], dtype=tf.float32, initializer=initializer)
        W0 = tf.get_variable("W0", [self.dim_o, self.dim_x], dtype=tf.float32, initializer=initializer)
        b0 = tf.get_variable("b0", [self.dim_o], dtype=tf.float32, initializer=initializer)
        # # 200 x 150 x 3
        # conv1 = tf.layers.conv2d(
        #     inputs=o,
        #     filters=dim_h1,
        #     kernel_size=(krn_h1, krn_h1),
        #     padding="valid",
        #     strides=(2,2),
        #     activation=tf.nn.relu,
        #     name="conv1"
        # )
        # # 98 x 73 x 16
        # pool1 = tf.layers.max_pooling2d(
        #     inputs=conv1,
        #     pool_size=(2,2),
        #     strides=2,
        #     name="pool1"
        # )
        # # 49 x 37 x 16
        # h1 = tf.reshape(pool1, [-1, dim_o])

        # Model input and output
        # h1 = tf.nn.relu(tf.matmul(o, W1) + b1)
        # s  = tf.add(tf.matmul(h1, W2), b2, name="s")
        self.s = tf.matmul(self.o + b0, W0, name="s")
        # s = tf.add(tf.matmul(h1, W0), b0, name="s")

    def add_loss_op(self):
        # Ground truth loss
        # loss = tf.losses.mean_squared_error(self.x, self.s)

        # Robotic rpiors loss
        L_temp = temporal_coherence_loss(self.s)
        L_prop = proportionality_loss(self.s, self.a)
        L_caus = causality_loss(self.s, self.a, self.r)
        L_rep  = repeatability_loss(self.s, self.a)
        self.loss = self.w_temp * L_temp + self.w_prop * L_prop + self.w_caus * L_caus + self.w_rep * L_rep

        tf.summary.scalar("loss", self.loss)

    def add_optimizer_op(self):
        # Create optimizer
        # optimizer = tf.train.AdadeltaOptimizer()
        # optimizer = tf.train.RMSPropOptimizer(0.0001)
        optimizer = tf.train.AdagradOptimizer(self.lr)
        self.train = optimizer.minimize(self.loss)

    def add_summary_op(self):
        self.summary = tf.summary.merge_all()
        self.saver = tf.train.Saver()

    def evaluate(self, observation):
        return self.sess.run(self.s, { self.o: observation })

    def create_logger(self):
        if not os.path.exists(os.path.join("results", "logs")):
            os.makedirs(os.path.join("results", "logs"))
        if not os.path.exists(os.path.join("results", "models")):
            os.makedirs(os.path.join("results", "models"))
        self.datadir = os.path.join("results", "logs", strftime("%m-%d_%H-%M"))
        self.modeldir = os.path.join("results", "models", strftime("%m-%d_%H-%M"))
        self.modelpath = os.path.join(self.modeldir, "model")
        if self.train_writer is not None:
            self.train_writer.close()
        self.train_writer = tf.summary.FileWriter(os.path.join(self.datadir, "train"))
        # test_writer = tf.summary.FileWriter(os.path.join(datadir, "test"))

    def reset_session(self):
        # Initialize session
        init = tf.global_variables_initializer()
        self.sess = tf.Session()
        self.sess.run(init) # reset values to wrong

    def train_iteration(self, o_train, a_train, r_train, x_train):
        feed_train = {
            self.o: o_train,
            self.x: x_train,
            self.a: a_train,
            self.r: r_train
        }
        self.sess.run(self.train, feed_train)
        summary_train, loss_train = self.sess.run([self.summary, self.loss], feed_train)

        return summary_train, loss_train

    def train_network(self, train_batch):
        min_loss_train = float("inf")

        i = 0
        for o_train, a_train, r_train, x_train, _, _ in train_batch:
            # Train iteration
            summary_train, loss_train = self.train_iteration(o_train, a_train, r_train, x_train)

            self.train_writer.add_summary(summary_train, i)

            if loss_train < min_loss_train:
                min_loss_train = loss_train
                self.saver.save(self.sess, self.modelpath, global_step=i)
                with open(os.path.join(self.modeldir, "saved.log"), "w+") as f:
                    f.write("Iteration: {}, Train loss: {}".format(i, loss_train))

            print("Iteration: {}, Train loss: {},".format(i, loss_train))
            i += 1

# Define custom py_func which takes also a grad op as argument:
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    # Need to generate a unique name to avoid duplicates:
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name, "PyFuncStateless": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

def temporal_coherence_loss(s, name=None):
    def _temporal_coherence_loss_grad(s):
        T = s.shape[0] - 1

        # Loss
        ds = s[1:] - s[:-1]

        loss = np.float32(1 / T) * (ds * ds).flatten().sum()

        # Gradient
        diff_ds = ds[:-1] - ds[1:]
        diff_ds = np.vstack((-ds[0], diff_ds, ds[-1]))

        # TODO: Should be T+1 x 2
        grad = 2 / T * diff_ds

        return loss, grad

    def _TemporalCoherenceLossGrad(op, grad_loss, grad_grad):
        return grad_loss * op.outputs[1]

    with ops.name_scope(name, "TemporalCoherenceLoss", [s]) as name:
        loss, _ = py_func(_temporal_coherence_loss_grad,
                          [s], [tf.float32, tf.float32], name=name,
                          grad=_TemporalCoherenceLossGrad)
        return loss

def sample_action_predicates(a):
    idx_a1 = []
    idx_a2 = []
    for i in range(a.shape[0]-1):
        idx_i_a  = np.linalg.norm(a[i,np.newaxis,:] - a[i+1:,:], axis=1) < k["epsilon_a"]
        idx_i_a  = np.random.permutation(np.where(idx_i_a)[0])
        num_a    = min(k["num_pair_samples"], len(idx_i_a))
        idx_i_a  = idx_i_a[:num_a] + i+1

        idx_a1.append(i + np.zeros(num_a, dtype=np.int32))
        idx_a2.append(idx_i_a)

    idx_a1 = np.concatenate(idx_a1)
    idx_a2 = np.concatenate(idx_a2)

    return (idx_a1, idx_a2)

def proportionality_loss(s, a, name=None):
    def _proportionality_loss_grad(s, a):
        T = s.shape[0] - 1
        S = s.shape[1]
        idx_a1, idx_a2 = sample_action_predicates(a)
        if idx_a1.shape[0] == 0:
            return np.float32(0), np.zeros(s.shape, dtype=np.float32)

        # Loss
        ds = s[1:] - s[:-1]
        ds_a1 = ds[idx_a1]
        ds_a2 = ds[idx_a2]

        ds_a1_norm = np.float32(np.linalg.norm(ds_a1, axis=1))
        ds_a2_norm = np.float32(np.linalg.norm(ds_a2, axis=1))

        # Prevent divide by 0
        idx = np.logical_and(ds_a1_norm != 0, ds_a2_norm != 0)
        if np.any(idx == 0):
            idx_a1 = idx_a1[idx]
            idx_a2 = idx_a2[idx]
            ds_a1 = ds[idx_a1]
            ds_a2 = ds[idx_a2]
            ds_a1_norm = np.float32(np.linalg.norm(ds_a1, axis=1))
            ds_a2_norm = np.float32(np.linalg.norm(ds_a2, axis=1))

        dds_norm = ds_a1_norm - ds_a2_norm

        loss = np.float32(1 / T) * dds_norm.dot(dds_norm)

        # Gradient
        I = np.eye(T+1, dtype=np.float32)
        D = I[1:] - I[:-1]
        diff_a1 = D[idx_a1]
        diff_a2 = D[idx_a2]

        q1 = diff_a1.T.dot((dds_norm / ds_a1_norm)[:,np.newaxis] * ds_a1)
        q2 = diff_a2.T.dot((dds_norm / ds_a2_norm)[:,np.newaxis] * ds_a2)

        grad = 2 / T * (q1 - q2)

        return loss, grad

    def _ProportionalityLossGrad(op, grad_loss, grad_grad):
        return grad_loss * op.outputs[1], 0 * op.inputs[1]

    with ops.name_scope(name, "ProportionalityLoss", [s, a]) as name:
        loss, _ = py_func(_proportionality_loss_grad,
                          [s, a], [tf.float32, tf.float32], name=name,
                          grad=_ProportionalityLossGrad)
        return loss

def sample_action_reward_predicates(a, r):
    idx_ar1 = []
    idx_ar2 = []
    for i in range(a.shape[0]-1):
        idx_i_a  = np.linalg.norm(a[i,np.newaxis,:] - a[i+1:,:], axis=1) < k["epsilon_a"]
        idx_i_r  = np.abs(r[i] - r[i+1:]) > k["epsilon_r"]
        idx_i_ar = np.logical_and(idx_i_a, idx_i_r)
        idx_i_ar = np.random.permutation(np.where(idx_i_ar)[0])
        num_ar   = min(k["num_pair_samples"], len(idx_i_ar))
        idx_i_ar = idx_i_ar[:num_ar] + i+1

        idx_ar1.append(i + np.zeros(num_ar, dtype=np.int32))
        idx_ar2.append(idx_i_ar)

    idx_ar1 = np.concatenate(idx_ar1)
    idx_ar2 = np.concatenate(idx_ar2)

    return (idx_ar1, idx_ar2)

def causality_loss(s, a, r, name=None):
    def _causality_loss_grad(s, a, r):
        T = s.shape[0] - 1
        S = s.shape[1]
        idx_ar1, idx_ar2 = sample_action_reward_predicates(a, r)
        if idx_ar1.shape[0] == 0:
            return np.float32(0), np.zeros(s.shape, dtype=np.float32)

        # Loss
        s = s[1:]
        s_ar1 = s[idx_ar1]
        s_ar2 = s[idx_ar2]
        ds = s_ar1 - s_ar2
        ds_norm = (ds * ds).sum(axis=1)

        c = ds_norm.min()
        exp_ds_norm = np.exp(c - ds_norm, dtype=np.float32)

        loss = np.float32(np.exp(-c) / T) * exp_ds_norm.sum()

        # Gradient
        I = np.eye(T+1, dtype=np.float32)
        d_ar = I[idx_ar1] - I[idx_ar2]
        d_ar_exp_ds_norm_ds = d_ar.T.dot(exp_ds_norm[:,np.newaxis] * ds)

        # TODO: Should be T+1 x 2
        grad = -2 * np.exp(-c) / T * d_ar_exp_ds_norm_ds

        return loss, grad

    def _CausalityLossGrad(op, grad_loss, grad_grad):
        return grad_loss * op.outputs[1], 0 * op.inputs[1], 0 * op.inputs[2]

    with ops.name_scope(name, "CausalityLoss", [s, a, r]) as name:
        loss, _ = py_func(_causality_loss_grad,
                          [s, a, r], [tf.float32, tf.float32], name=name,
                          grad=_CausalityLossGrad)
        return loss

def repeatability_loss(s, a, name=None):
    def _repeatability_loss_grad(s, a):
        T = s.shape[0] - 1
        S = s.shape[1]
        idx_a1, idx_a2 = sample_action_predicates(a)
        if idx_a1.shape[0] == 0:
            return np.float32(0), np.zeros(s.shape, dtype=np.float32)

        # Loss
        ds = s[1:] - s[:-1]
        s  = s[1:]
        s_a1  = s[idx_a1]
        s_a2  = s[idx_a2]
        ds_a1 = ds[idx_a1]
        ds_a2 = ds[idx_a2]
        ds  = s_a1 - s_a2
        dds = ds_a1 - ds_a2
        ds_norm  = (ds * ds).sum(axis=1)
        dds_norm = (dds * dds).sum(axis=1)
        c = ds_norm.min()
        exp_ds_norm = np.exp(c - ds_norm)
        exp_ds_dds_norm = exp_ds_norm * dds_norm

        loss = np.float32(np.exp(-c) / T) * exp_ds_dds_norm.sum()

        # Gradient
        # TODO: Find more efficient way to find A1 - A2
        I = np.eye(T, dtype=np.float32)
        da = I[idx_a1] - I[idx_a2]
        I = np.eye(T+1, dtype=np.float32)
        D  = I[1:] - I[:-1]
        diff_da = da.dot(D)

        X = diff_da.T.dot(exp_ds_norm[:,np.newaxis] * dds)
        Y = da.T.dot(exp_ds_dds_norm[:,np.newaxis] * ds)
        Y = np.vstack((np.zeros(Y.shape[1], dtype=np.float32), Y))

        grad = np.float32(2 * np.exp(-c) / T) * (X - Y)

        return loss, grad

    def _RepeatabilityLossGrad(op, grad_loss, grad_grad):
        return grad_loss * op.outputs[1], 0 * op.inputs[1]

    with ops.name_scope(name, "RepeatabilityLoss", [s, a]) as name:
        loss, _ = py_func(_repeatability_loss_grad,
                             [s, a], [tf.float32, tf.float32], name=name,
                             grad=_RepeatabilityLossGrad)
        return loss

<file_sep>import gym
from gym import error, spaces, utils
from gym.utils import seeding
import numpy as np

from gym_sai2.envs import sai2_env_cpp
 
class SaiEnv(gym.Env):

    def __init__(self):
        world_file    = "../resources/gym.urdf"
        robot_file    = "../resources/kuka_iiwa_gym.urdf"
        robot_name    = "kuka_iiwa"
        window_width  = 200
        window_height = 150

        self.dim_action      = (2,)
        self.dim_observation = (window_height, window_width, 3)
        self.dim_x           = (2,)

        self.action_space      = spaces.Box(-1, 1, self.dim_action)
        self.observation_space = spaces.Box(0, 255, self.dim_observation)

        self.img_buffer  = np.zeros(self.dim_observation, dtype=np.uint8)
        self.info_buffer = np.zeros((2,3), dtype=np.float64)
        self.sai2_env    = sai2_env_cpp.init(world_file, robot_file, robot_name,
                                             window_width, window_height)

    def _step(self, action):
        """
        Run one timestep of the environment's dynamics. When end of episode is
        reached, you are responsible for calling `reset()` to reset this
        environment's state.
        
        Accepts an action and returns a tuple (observation, reward, done, info).
        
        Args:
            action (object): an action provided by the environment
            
        Returns:
            observation (object): agent's observation of the current environment
            reward (float): amount of reward returned after previous action
            done (boolean): whether the episode has ended, in which case further step() calls will return undefined results
            info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
        """
        reward, done = sai2_env_cpp.step(self.sai2_env, action, self.img_buffer, self.info_buffer)
        info = {
            "x": self.info_buffer[0],
            "dx": self.info_buffer[1],
        }
        return (self.img_buffer, reward, done, info)

    def _reset(self):
        """
        Resets the state of the environment and returns an initial observation.

        Returns:
            observation (object): the initial observation of the space.
        """
        sai2_env_cpp.reset(self.sai2_env, self.img_buffer)
        return self.img_buffer

    def _render(self, mode="human", close=False):
        """
        Renders the environment. The set of supported modes varies per
        environment. (And some environments do not support rendering at all.) By
        convention, if mode is:

        - human: render to the current display or terminal and return nothing.
          Usually for human consumption.
        - rgb_array: Return an numpy.ndarray with shape (x, y, 3), representing
          RGB values for an x-by-y pixel image, suitable for turning into a
          video.
        - ansi: Return a string (str) or StringIO.StringIO containing a
          terminal-style text representation. The text can include newlines and
          ANSI escape sequences (e.g. for colors).

        Note:
            Make sure that your class's metadata 'render.modes' key includes the
            list of supported modes. It's recommended to call super() in
            implementations to use the functionality of this method.
            
        Args:
            mode (str): the mode to render with
            close (bool): close all open renderings

        Example:

        class MyEnv(Env):
            metadata = {'render.modes': ['human', 'rgb_array']}

            def render(self, mode='human'):
                if mode == 'rgb_array':
                    return np.array(...) # return RGB frame suitable for video
                elif mode is 'human':
                  ... # pop up a window and render
                else:
                  super(MyEnv, self).render(mode=mode) # just raise an exception
        """
        if mode == "rgb_array":
            return self.img_buffer
        return

    def _seed(self, seed=None):
        """
        Sets the seed for this env's random number generator(s).

        Note:
            Some environments use multiple pseudorandom number
            generators. We want to capture all such seeds used in order to
            ensure that there aren't accidental correlations between multiple
            generators.

        Returns:
            list<bigint>: Returns the list of seeds used in this env's random
            number generators. The first value in the list should be the "main"
            seed, or the value which a reproducer should pass to 'seed'. Often,
            the main seed equals the provided 'seed', but this won't be true if
            seed=None, for example.
        """
        self.np_random, seed = seeding.np_random(seed)
        return [seed]
<file_sep>#!/usr/bin/env python

import gym
import gym_sai2
import threading

import numpy as np
import matplotlib.pyplot as plt
import h5py
from time import gmtime, strftime

def wait():
    input("Hit <enter> to close.\n")
    exit()

class RandomAgent(object):
    def __init__(self, action_space):
        self.action_space = action_space

    def act(self, observation, reward, done):
        return self.action_space.sample()

if __name__ == "__main__":
    NUM_BATCHES = 100
    SIZE_BATCH  = 1000

    # Create thread to kill process (ctrl-c doesn't work?)
    thread = threading.Thread(target=wait, daemon=True)
    thread.start()

    # Create sai2 environment
    env = gym.make("sai2-v0")
    env.seed(0)
    agent = RandomAgent(env.action_space)

    # Create dataset
    filename = "data/data-{}.hdf5".format(strftime("%m-%d_%H-%M"), gmtime())
    with h5py.File(filename, "w") as f:

        # Get initial observation
        initial_observation = env.reset()
        ob = initial_observation
        reward = 0
        done = False
        dset = f.create_dataset("initial_observation", initial_observation.shape, dtype=initial_observation.dtype)
        dset[...] = initial_observation

        for i in range(NUM_BATCHES):

            print("Iteration: {}".format(i))

            actions = []
            observations = []
            rewards = []
            xs = []
            dxs = []

            # Generate trajectory
            for _ in range(SIZE_BATCH):
                action = agent.act(ob, reward, done)
                ob, reward, done, info = env.step(action)

                actions.append(np.array(action))
                observations.append(np.array(ob)[np.newaxis,...])
                rewards.append(reward)
                xs.append(np.array(info["x"]))
                dxs.append(np.array(info["dx"]))

            # Vectorize trajectory
            actions = np.row_stack(actions)
            observations = np.concatenate(observations, axis=0)
            rewards = np.array(rewards)
            xs = np.row_stack(xs)
            dxs = np.row_stack(dxs)

            # Save trajectory to dataset
            grp = f.create_group("/episodes/{0:05d}".format(i))
            dset = grp.create_dataset("actions", actions.shape, dtype=actions.dtype)
            dset[...] = actions
            dset = grp.create_dataset("observations", observations.shape, dtype=observations.dtype)
            dset[...] = observations
            dset = grp.create_dataset("rewards", rewards.shape, dtype=rewards.dtype)
            dset[...] = rewards
            dset = grp.create_dataset("xs", xs.shape, dtype=xs.dtype)
            dset[...] = xs
            dset = grp.create_dataset("dxs", dxs.shape, dtype=dxs.dtype)
            dset[...] = dxs

        env.close()

<file_sep>#!/usr/bin/env python

import gym
import gym_sai2

import numpy as np
from data import *
from robotic_priors import *
from reinforcement_learning import *
import threading

def wait():
    input("Hit <enter> to close.\n")
    exit()

if __name__ == "__main__":
    NUM_EPISODES = 10
    LEN_EPISODE  = 1000
    NUM_ITERATIONS = 20

    # Create thread to kill process (ctrl-c doesn't work?)
    thread = threading.Thread(target=wait, daemon=True)
    thread.start()

    # Create sai2 environment
    env = gym.make("sai2-v0")
    # TODO: Implement environment seed
    env.seed(0)

    # State representation learning
    robotic_priors = RoboticPriors(env.dim_observation)
    robotic_priors.reset_session()

    # Reinforcement learning agent
    agent = RandomAgent(env.action_space)

    for i in range(NUM_ITERATIONS):

        print("Iteration: {}".format(i))

        # Generate RL trajectories
        with DataLogger() as d:
            filename = d.filename

            for j in range(NUM_EPISODES):

                print("Episode: {}".format(j))

                # Get initial observation
                ob = env.reset()
                ob = np.array(ob)[np.newaxis,...]

                # One episode
                actions, observations, rewards, xs, dxs, learned_states = [], [ob], [], [], [], []

                # Generate trajectory
                for _ in range(LEN_EPISODE):
                    # Query RL policy
                    s_hat = robotic_priors.evaluate(np.reshape(ob, (1,-1)))
                    action = agent.act(s_hat)
                    ob, reward, done, info = env.step(action)
                    ob = np.array(ob)[np.newaxis,...]

                    # Append to trajectory
                    actions.append(np.array(action))
                    observations.append(ob)
                    rewards.append(reward)
                    xs.append(np.array(info["x"]))
                    dxs.append(np.array(info["dx"]))
                    learned_states.append(s_hat)

                # Vectorize trajectory
                actions = np.row_stack(actions)
                observations = np.concatenate(observations, axis=0)
                rewards = np.array(rewards)
                xs = np.row_stack(xs)
                dxs = np.row_stack(dxs)
                learned_states = np.row_stack(learned_states)

                # Log data
                d.log(j, actions, observations, rewards, xs, dxs, learned_states)

            # Get list of actions, observations, rewards, xs, dxs, learned_states from all episodes
            episodes = d.flush()

        # Train represnetation learning
        robotic_priors.create_logger()
        robotic_priors_data_generator = batch_data(data=episodes, extra=True, flatten=True)
        robotic_priors.train_network(robotic_priors_data_generator)

    env.close()

<file_sep>set -e

mkdir -p build
cd build
cmake ..
make -j4
cd ..

# Insert helper scripts into bin directory
cd bin

# Make script
cat <<EOF > make.sh
cd ..
mkdir -p build
cd build
cmake ..
make -j4
cd ../bin
EOF
chmod +x make.sh

# Run generic controller script
cat <<EOF > visualizer.sh
# /opt/VirtualGL/bin/vglrun ./sai2_env ../resources/gym.urdf ../resources/kuka_iiwa_gym.urdf kuka_iiwa
./sai2_env ../resources/gym.urdf ../resources/kuka_iiwa_gym.urdf kuka_iiwa
EOF
chmod +x visualizer.sh

cd ..
<file_sep># Set executable name
set(EXECUTABLE_NAME sai2_env)

# Create executable
add_executable(${EXECUTABLE_NAME}
	"${CMAKE_CURRENT_SOURCE_DIR}/CSaiEnv.cpp"
)

# Link libraries
find_library(GLFW_LIBRARY glfw)
target_link_libraries(${EXECUTABLE_NAME}
	${SAI2_LIBRARIES}
	${GLFW_LIBRARY}
)

# Static library for python interface
set(LIBRARY_NAME sai2-env)

# Create library
add_library(${LIBRARY_NAME} SHARED
	"${CMAKE_CURRENT_SOURCE_DIR}/sai2_env_py.cpp"
	"${CMAKE_CURRENT_SOURCE_DIR}/CSaiEnv.cpp"
)

# Link libraries
find_package(PythonLibs)
include_directories(${PYTHON_INCLUDE_DIRS})
target_link_libraries(${LIBRARY_NAME}
	${SAI2_LIBRARIES}
	${GLFW_LIBRARY}
	${PYTHON_LIBRARIES}
)

<file_sep>#include <Python.h>

#include <string>
#include <iostream>

#include "CSaiEnv.h"

extern "C" {
	void *init(const char *cstr_world_file, const char *cstr_robot_file, const char *cstr_robot_name,
	           size_t window_width, size_t window_height);
	bool step(void *p_sai2_env, const double *action, uint8_t *observation, double *reward, double *info);
	void reset(void *p_sai2_env, uint8_t *observation);
}

void *init(const char *cstr_world_file, const char *cstr_robot_file, const char *cstr_robot_name,
           size_t window_width, size_t window_height) {
	std::string world_file(cstr_world_file);
	std::string robot_file(cstr_robot_file);
	std::string robot_name(cstr_robot_name);

	// Load robot
	std::cout << "Loading robot: " << robot_file << std::endl;
	auto robot = std::make_shared<Model::ModelInterface>(robot_file, Model::rbdl, Model::urdf, false);
	auto sim = std::make_shared<Simulation::SimulationInterface>(world_file, Simulation::sai2simulation, Simulation::urdf, false);
	auto graphics = std::make_shared<Graphics::GraphicsInterface>(world_file, Graphics::chai, Graphics::urdf, true);

	// Start controller app
	std::cout << "Initializing app with " << robot_name << std::endl;
	CSaiEnv *sai2_env = new CSaiEnv(std::move(robot), std::move(sim), std::move(graphics), robot_name, window_width, window_height);
	sai2_env->initialize();

	return sai2_env;
}

bool step(void *p_sai2_env, const double *action, uint8_t *observation, double *reward, double *info) {
	CSaiEnv *sai2_env = reinterpret_cast<CSaiEnv *>(p_sai2_env);
	return sai2_env->step(action, observation, *reward, info);
}

void reset(void *p_sai2_env, uint8_t *observation) {
	CSaiEnv *sai2_env = reinterpret_cast<CSaiEnv *>(p_sai2_env);
	sai2_env->reset(observation);
}
<file_sep>gym==0.9.5
h5py==2.7.1
jupyter==1.0.0
matplotlib==2.1.2
tensorflow==2.7.2
<file_sep>#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from data import *
import os
from time import strftime
from robotic_priors import *

# Define robotic prior constants
k = {
    "epsilon_a": 0.01,
    "epsilon_r": 0.01,
    "range_W": 0.05,
    "w_temp": 5,
    "w_prop": 1,
    "w_caus": 5,
    "w_rep": 5,
    "lambda": 0.0025,
    "dim_s": 2,
    "num_pair_samples": 1000
}

if __name__ == "__main__":
    LOG = True

    print(list_files())
    filename = get_filename(-1)
    print(filename)

    # Model parameters
    dim_image = np.array(get_observation_dim(filename))
    dim_o = np.prod(dim_image)
    dim_x = 2
    dim_a = 2
    dim_h1 = 16
    dim_h2 = 8
    krn_h1 = 5
    krn_h2 = 5

    # Initialize train and test datasets
    train_batch = batch_data(size_batch=1000, extra=True, filename=filename, dataset="train", flatten=True)
    test_batch  = batch_data(size_batch=1000, extra=True, filename=filename, dataset="test", flatten=True)
    o_test, a_test, r_test, x_test,  _ = next(test_batch)

    # Placeholders
    o = tf.placeholder(tf.float32, [None, dim_o], "o")
    a = tf.placeholder(tf.float32, [None, dim_a], "a")
    r = tf.placeholder(tf.float32, [None], "r")
    x = tf.placeholder(tf.float32, [None, dim_x], "x")
    feed_test = {
        o: o_test,
        x: x_test,
        a: a_test,
        r: r_test
    }

    # Model architecture
    initializer = tf.contrib.layers.xavier_initializer()
    # W1 = tf.get_variable("W1", [dim_o, dim_h1], dtype=tf.float32, initializer=initializer)
    # W2 = tf.get_variable("W2", [dim_h1, dim_x], dtype=tf.float32, initializer=initializer)
    # b1 = tf.get_variable("b1", [dim_h1], dtype=tf.float32, initializer=initializer)
    # b2 = tf.get_variable("b2", [dim_x], dtype=tf.float32, initializer=initializer)
    W0 = tf.get_variable("W0", [dim_o, dim_x], dtype=tf.float32, initializer=initializer)
    b0 = tf.get_variable("b0", [dim_o], dtype=tf.float32, initializer=initializer)
    # # 200 x 150 x 3
    # conv1 = tf.layers.conv2d(
    #     inputs=o,
    #     filters=dim_h1,
    #     kernel_size=(krn_h1, krn_h1),
    #     padding="valid",
    #     strides=(2,2),
    #     activation=tf.nn.relu,
    #     name="conv1"
    # )
    # # 98 x 73 x 16
    # pool1 = tf.layers.max_pooling2d(
    #     inputs=conv1,
    #     pool_size=(2,2),
    #     strides=2,
    #     name="pool1"
    # )
    # # 49 x 37 x 16
    # h1 = tf.reshape(pool1, [-1, dim_o])

    # Model input and output
    # h1 = tf.nn.relu(tf.matmul(o, W1) + b1)
    # s  = tf.add(tf.matmul(h1, W2), b2, name="s")
    s = tf.matmul(o + b0, W0, name="s")
    # s = tf.add(tf.matmul(h1, W0), b0, name="s")

    # Ground truth loss
    # loss = tf.losses.mean_squared_error(x, s)

    # Robotic rpiors loss
    L_temp = temporal_coherence_loss(s)
    L_prop = proportionality_loss(s, a)
    L_caus = causality_loss(s, a, r)
    L_rep  = repeatability_loss(s, a)
    loss = k["w_temp"] * L_temp + k["w_prop"] * L_prop + k["w_caus"] * L_caus + k["w_rep"] * L_rep

    tf.summary.scalar("loss", loss)

    # Create optimizer
    # optimizer = tf.train.AdadeltaOptimizer()
    # optimizer = tf.train.RMSPropOptimizer(0.0001)
    optimizer = tf.train.AdagradOptimizer(0.001)
    train = optimizer.minimize(loss)

    # Create logger
    merged = tf.summary.merge_all()
    if LOG:
        datadir = os.path.join("data", "logs", strftime("%m-%d_%H-%M"))
        modeldir = os.path.join("data", "models", strftime("%m-%d_%H-%M"))
        modelpath = os.path.join(modeldir, "model")
        train_writer = tf.summary.FileWriter(os.path.join(datadir, "train"))
        test_writer = tf.summary.FileWriter(os.path.join(datadir, "test"))
        saver = tf.train.Saver()
        min_loss_test = float("inf")

    # Initialize session
    init = tf.global_variables_initializer()
    sess = tf.Session()
    sess.run(init) # reset values to wrong
    for i in range(10000):
        # Train iteration
        o_train, a_train, r_train, x_train, _ = next(train_batch)
        feed_train = {
            o: o_train,
            x: x_train,
            a: a_train,
            r: r_train
        }
        sess.run(train, feed_train)

        # Evaluate training accuracy
        summary_train, loss_train = sess.run([merged, loss], feed_train)
        summary_test, loss_test  = sess.run([merged, loss], feed_test)

        # Log summary
        if LOG and i % 100 == 0:
            train_writer.add_summary(summary_train, i)
            test_writer.add_summary(summary_test, i)
        if LOG:
            if loss_test < min_loss_test:
                min_loss_test = loss_test
                saver.save(sess, modelpath, global_step=i)
                with open(os.path.join(modeldir, "saved.log"), "w+") as f:
                    f.write("Iteration: {}, Train loss: {}, Test loss: {}".format(i, loss_train, loss_test))

        print("Iteration: {}, Train loss: {}, Test loss: {}".format(i, loss_train, loss_test))

    train_writer.close()
    test_writer.close()
<file_sep>find_program(PYTHON "python")
if (PYTHON)
	set(SETUP_PY_IN "${CMAKE_CURRENT_SOURCE_DIR}/setup.py.in")
	set(SETUP_PY    "${CMAKE_CURRENT_BINARY_DIR}/setup.py")
	set(DEPS        "${CMAKE_CURRENT_SOURCE_DIR}/gym_sai2/__init__.py")
	set(OUTPUT      "${CMAKE_CURRENT_SOURCE_DIR}/build/timestamp")

	configure_file(${SETUP_PY_IN} ${SETUP_PY})

	add_custom_command(OUTPUT ${OUTPUT}
	    COMMAND ${PYTHON} ${SETUP_PY} build
		COMMAND ${CMAKE_COMMAND} -E touch ${OUTPUT}
		DEPENDS ${DEPS})

	add_custom_target(target ALL DEPENDS ${OUTPUT})

	install(CODE "execute_process(COMMAND ${PYTHON} ${SETUP_PY} install)")
endif()
<file_sep>cmake_minimum_required(VERSION 3.1)

project(SAI2-MASTER)

# CMake settings
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)


######################
# External libraries #
######################

set(SAI2_EXTERNAL_DIR ${PROJECT_SOURCE_DIR}/lib)

# Eigen
set(EIGEN3_INCLUDE_DIR ${SAI2_EXTERNAL_DIR}/eigen)

# RBDL
# TODO use FindRBDL.cmake in rbdl/examples
set(RBDL_DIR ${SAI2_EXTERNAL_DIR}/rbdl)
file(GLOB RBDL_PATHS ${RBDL_DIR}/build* ${RBDL_DIR}/build*/addons/urdfreader)
find_library(RBDL_LIBRARY rbdl REQUIRED PATHS ${RBDL_PATHS})
find_library(RBDL_URDFReader_LIBRARY rbdl_urdfreader REQUIRED PATHS ${RBDL_PATHS})
find_path(RBDL_URDFReader_INCLUDE_DIR rbdl/addons/urdfreader/urdfreader.h HINTS ${SAI2_EXTERNAL_DIR})
file(GLOB RBDL_INCLUDE_DIR ${RBDL_DIR}/include ${RBDL_DIR}/build*/include)
set(RBDL_FOUND true)

# Chai3d
# TODO: Find better place to put this?
find_package(CHAI3D PATHS ${SAI2_EXTERNAL_DIR}/chai3d.git/build)
find_library(DRD_LIBRARY drd PATHS ${CHAI3D_LIBRARY_DIRS})
find_library(CHAI3D_LIBRARY chai3d paths ${CHAI3D_DIR})
list(REMOVE_ITEM CHAI3D_LIBRARIES drd)
list(REMOVE_ITEM CHAI3D_LIBRARIES chai3d)
set(CHAI3D_LIBRARIES ${CHAI3D_LIBRARY} ${DRD_LIBRARY} ${CHAI3D_LIBRARIES})

# SAI2 Common
set(BUILD_EXAMPLES OFF CACHE BOOL "")
add_subdirectory(${SAI2_EXTERNAL_DIR}/sai2-simulation.git)
add_subdirectory(${SAI2_EXTERNAL_DIR}/sai2-common.git)

# TODO: Find local way to do this
find_package(SAI2-COMMON PATHS ${SAI2_EXTERNAL_DIR}/sai2-common.git/build)
include_directories(${SAI2-COMMON_INCLUDE_DIRS})
message(${SAI2-COMMON_INCLUDE_DIRS})


###################
# Project Sources #
###################

# Set include include directories
include_directories(${SAI2_EXTERNAL_DIR})
include_directories(${PROJECT_SOURCE_DIR}/src)

# Set common libraries to be linked
set(SAI2_LIBRARIES
	${SAI2-COMMON_LIBRARY}
)

# Add specific apps
add_subdirectory(src/gym_sai2)
# add_subdirectory(python)
<file_sep>from gym.envs.registration import register

register(
    id="sai2-v0",
    entry_point="gym_sai2.envs:SaiEnv"
)
<file_sep>import ctypes as c

sai2py = c.cdll.LoadLibrary("./libsai2-env.so")

c_ubyte_p = c.POINTER(c.c_ubyte)
c_double_p = c.POINTER(c.c_double)

sai2_init = sai2py.init
sai2_init.restype = c.c_void_p
sai2_init.argtypes = [c.c_char_p, c.c_char_p, c.c_char_p, c.c_size_t, c.c_size_t]

sai2_step = sai2py.step
sai2_step.restype = c.c_bool
sai2_step.argtypes = [c.c_void_p, c_double_p, c_ubyte_p, c_double_p, c_double_p]

sai2_reset = sai2py.reset
sai2_reset.restype = None
sai2_reset.argtypes = [c.c_void_p, c_ubyte_p]

def init(world_file, robot_file, robot_name, window_width, window_height):
    return sai2_init(c.c_char_p(world_file.encode()),
                     c.c_char_p(robot_file.encode()),
                     c.c_char_p(robot_name.encode()),
                     c.c_size_t(window_width),
                     c.c_size_t(window_height))

def step(sai2_env, action, observation, info):
    reward = c.c_double()
    done = sai2_step(sai2_env,
                     action.ctypes.data_as(c_double_p),
                     observation.ctypes.data_as(c_ubyte_p),
                     c.byref(reward),
                     info.ctypes.data_as(c_double_p))
    return (reward.value, done)

def reset(sai2_env, observation):
    sai2_reset(sai2_env,
               observation.ctypes.data_as(c_ubyte_p))
<file_sep>#ifndef C_SAI_ENV_H
#define C_SAI_ENV_H

// CS225a
#include <redis/RedisClient.h>
#include <timer/LoopTimer.h>

// Standard
#include <string>
#include <thread>
#include <mutex>
#include <condition_variable>

// External
#include <Eigen/Core>
#include <model/ModelInterface.h>
#include <simulation/SimulationInterface.h>
#include <graphics/GraphicsInterface.h>
#include <graphics/ChaiGraphics.h>
#include <GLFW/glfw3.h> //must be loaded after loading opengl/glew


class CSaiEnv {

public:

	CSaiEnv(std::shared_ptr<Model::ModelInterface> robot,
	        std::shared_ptr<Simulation::SimulationInterface> sim,
	        std::shared_ptr<Graphics::GraphicsInterface> graphics,
	        std::string& robot_name,
	        size_t window_width = 400,
	        size_t window_height = 300) :
		robot_(robot),
		dof(robot->dof()),
		kRobotName(robot_name),
		sim(sim),
		kWindowWidth(window_width),
		kWindowHeight(window_height),
		thread_graphics(&CSaiEnv::graphicsMain, this, graphics),
		// graphics(dynamic_cast<Graphics::ChaiGraphics *>(graphics->_graphics_internal)),
		KEY_COMMAND_TORQUES (RedisServer::KEY_PREFIX + robot_name + "::actuators::fgc"),
		KEY_EE_POS          (RedisServer::KEY_PREFIX + robot_name + "::tasks::ee_pos"),
		KEY_EE_POS_DES      (RedisServer::KEY_PREFIX + robot_name + "::tasks::ee_pos_des"),
		KEY_JOINT_POSITIONS (RedisServer::KEY_PREFIX + robot_name + "::sensors::q"),
		KEY_JOINT_VELOCITIES(RedisServer::KEY_PREFIX + robot_name + "::sensors::dq"),
		KEY_TIMESTAMP       (RedisServer::KEY_PREFIX + robot_name + "::timestamp"),
		KEY_KP_POSITION     (RedisServer::KEY_PREFIX + robot_name + "::tasks::kp_pos"),
		KEY_KV_POSITION     (RedisServer::KEY_PREFIX + robot_name + "::tasks::kv_pos"),
		KEY_KP_ORIENTATION  (RedisServer::KEY_PREFIX + robot_name + "::tasks::kp_ori"),
		KEY_KV_ORIENTATION  (RedisServer::KEY_PREFIX + robot_name + "::tasks::kv_ori"),
		KEY_KP_JOINT        (RedisServer::KEY_PREFIX + robot_name + "::tasks::kp_joint"),
		KEY_KV_JOINT        (RedisServer::KEY_PREFIX + robot_name + "::tasks::kv_joint"),
		command_torques_(dof),
		q_des_(dof),
		dq_des_(dof)
	{
		reset();
	}

	/***** Public functions *****/

	void initialize();
	void runLoop();
	bool step(const double *action, uint8_t *observation, double& reward, double *info);
	void reset(uint8_t *observation);

protected:

	/***** Enums *****/

	// State enum for controller state machine inside runloop()
	enum ControllerState {
		REDIS_SYNCHRONIZATION,
		JOINT_SPACE_INITIALIZATION,
		OP_SPACE_POSITION_CONTROL
	};

	// Return values from computeControlTorques() methods
	enum ControllerStatus {
		RUNNING,  // Not yet converged to goal position
		FINISHED  // Converged to goal position
	};

	/***** Constants *****/

	const int dof;  // Initialized with robot model
	const double kToleranceInitQ  = 0.1;  // Joint space initialization tolerance
	const double kToleranceInitDq = 0.1;  // Joint space initialization tolerance
	const double kMaxVelocity = 0.5;  // Maximum end effector velocity
	const double kToleranceInitX  = 0.01;  // Task space initialization tolerance

	const int kControlFreq = 1000;         // 1 kHz control loop
	const int kSimulationFreq = 10000;         // 1 kHz control loop
	const int kEnvironmentFreq = 100;         // 1 kHz control loop
	const int kInitializationPause = 1e6;  // 1ms pause before starting control loop

	const std::string kRedisHostname = "127.0.0.1";
	const int kRedisPort = 6379;

	const std::string kRobotName;

	const std::string kCameraName = "camera_fixed";
	const Eigen::Vector3d kCameraPos = Eigen::Vector3d(-0.5, 0.0, 1.1);
	const Eigen::Vector3d kCameraVertical = Eigen::Vector3d(0, 0, 1);
	const Eigen::Vector3d kCameraLookat = Eigen::Vector3d(0, -0.3, 0.5);
	// const Eigen::Vector3d kCameraPos = Eigen::Vector3d(-0.8, -0.1, 0.9);
	// const Eigen::Vector3d kCameraVertical = Eigen::Vector3d(0, 0, 1);
	// const Eigen::Vector3d kCameraLookat = Eigen::Vector3d(0, -0.3, 0.5);
	const int kWindowWidth;
	const int kWindowHeight;

	const double kWallTolerance = 0.02;
	const double kCornerDistance = 0.1;
	const Eigen::Vector3d kCenter = Eigen::Vector3d(0, -0.45, 0.55);
	const double kCenterDistance = 0.2;
	const double kTimeEpisode = 10;

	// Redis keys:
	// - write:
	const std::string KEY_COMMAND_TORQUES;
	const std::string KEY_EE_POS;
	const std::string KEY_EE_POS_DES;
	// - read:
	const std::string KEY_JOINT_POSITIONS;
	const std::string KEY_JOINT_VELOCITIES;
	const std::string KEY_TIMESTAMP;
	const std::string KEY_KP_POSITION;
	const std::string KEY_KV_POSITION;
	const std::string KEY_KP_ORIENTATION;
	const std::string KEY_KV_ORIENTATION;
	const std::string KEY_KP_JOINT;
	const std::string KEY_KV_JOINT;

	/***** Member functions *****/

	void reset();
	void readRedisValues();
	void updateModel();
	void writeRedisValues();
	ControllerStatus computeJointSpaceControlTorques();
	ControllerStatus computeOperationalSpaceControlTorques();
	ControllerStatus computeSimpleOperationalSpaceControlTorques();

	void syncGraphics();
	void graphicsMain(std::shared_ptr<Graphics::GraphicsInterface> graphics);

	/***** Member variables *****/

	// Robot
	const std::shared_ptr<Model::ModelInterface> robot_;
	const std::shared_ptr<Simulation::SimulationInterface> sim;
	Graphics::ChaiGraphics *graphics;

	bool update_graphics_ = false;
	std::mutex mutex_graphics_;
	std::mutex mutex_robot_;
	std::condition_variable cv_;
	std::thread thread_simulator;
	std::thread thread_graphics;

	// Redis
	RedisClient redis_;

	// Timer
	LoopTimer timer_;
	double t_curr_;
	uint64_t controller_counter_ = 0;

	// State machine
	ControllerState controller_state_ = REDIS_SYNCHRONIZATION;

	// Controller variables
	Eigen::VectorXd command_torques_;
	Eigen::MatrixXd J_;
	Eigen::MatrixXd N_;
	Eigen::MatrixXd Lambda_ = Eigen::MatrixXd(6, 6);
	Eigen::VectorXd g_;
	Eigen::Vector3d x_, dx_;
	Eigen::Matrix3d R_des_, R_ee_to_base_;
	Eigen::Vector3d w_;
	Eigen::VectorXd q_des_, dq_des_;
	Eigen::Vector3d x_des_, dx_des_;
	Eigen::Vector3d action_;

	// Graphics
	GLubyte *gl_buffer_ = nullptr;

	// Default gains (used only when keys are nonexistent in Redis)
	double kp_action_ = 100;
	double kp_pos_ = 100;
	double kv_pos_ = 20;
	double kp_ori_ = 200;
	double kv_ori_ = 40;
	double kp_joint_ = 40;
	double kv_joint_ = 10;
};

#endif  // C_SAI_ENV_H
<file_sep># cd NatNetLinux.git
# git submodule update --init
# cd ..

# TODO: Place inside sai2-common
curl -L https://bitbucket.org/eigen/eigen/get/3.3.4.zip -o eigen-3.3.4.zip
unzip eigen-3.3.4.zip
mv eigen-eigen-* eigen

# TODO: Place inside sai2-common
curl -L https://bitbucket.org/rbdl/rbdl/get/v2.5.0.zip -o rbdl-2.5.0.zip
unzip rbdl-2.5.0.zip
mv rbdl-rbdl-* rbdl
cd rbdl
mkdir -p build
cd build
cmake -DCMAKE_BUILD_TYPE=Release -DRBDL_BUILD_ADDON_URDFREADER=ON -DRBDL_USE_ROS_URDF_LIBRARY=OFF ..
make -j4
cd ../..

cd chai3d.git
git submodule update --init
mkdir -p build
cd build
cmake -DCMAKE_BUILD_TYPE=Release ..
make -j4
cd ../..

# cd sai2-common.git
# # git submodule update --init --recursive
# mkdir -p build_rel
# cd build_rel
# cmake -DCMAKE_BUILD_TYPE=Release ..
# make -j4
# cd ../..

# cd sai2-simulation.git
# # git submodule update --init --recursive
# mkdir -p build_rel
# cd build_rel
# cmake -DCMAKE_BUILD_TYPE=Release ..
# make -j4
# cd ../..

<file_sep>#include "CSaiEnv.h"

#include <iostream>
#include <fstream>
#include <cmath>

#include <signal.h>
#include <stdlib.h>
static volatile bool g_runloop = true;
void stop(int) { g_runloop = false; }

// Return true if any elements in the Eigen::MatrixXd are NaN
template<typename Derived>
static inline bool isnan(const Eigen::MatrixBase<Derived>& x) {
	return (x.array() != x.array()).any();
}

static void glfwError(int error, const char* description) {
	std::cerr << "GLFW Error: " << description << std::endl;
	exit(1);
}

/**
 * CSaiEnv::reset()
 * ----------------
 * Reset robot simulation.
 */
void CSaiEnv::reset() {
	controller_counter_ = 0;
	command_torques_.setZero();

	// Home configuration for Kuka iiwa
	// q_des_ << 90, -30, 0, 60, 0, -90, -60;
	// q_des_ *= M_PI / 180.0;
	q_des_ << 1.570796, -0.151874, 0, 1.681582, 0, -1.308137, -1.570796;
	dq_des_.setZero();

	// Desired end effector position
	x_des_ << kCenterDistance * (2 * (double)rand() / RAND_MAX - 1) + kCenter(0),
	          kCenterDistance * (2 * (double)rand() / RAND_MAX - 1) + kCenter(1),
	          kCenter(2);
	action_ << 0, 0, 0;
	dx_des_.setZero();
	R_des_ << 1,  0,  0,
		      0, -1,  0,
		      0,  0, -1;

	// Reset model
	sim->setJointPositions(kRobotName, q_des_);
	sim->setJointVelocities(kRobotName, dq_des_);
	try {
		readRedisValues();
	} catch (std::exception& e) {
		std::cout << e.what() << std::endl;
	}
	updateModel();

	// 10s to get to random x_des
	while ((x_ - x_des_).norm() > kToleranceInitX) {
		computeSimpleOperationalSpaceControlTorques();
		sim->setJointTorques(kRobotName, command_torques_);
		sim->integrate(1.0 / kControlFreq);
		sim->getJointPositions(kRobotName, robot_->_q);
		sim->getJointVelocities(kRobotName, robot_->_dq);
		updateModel();
	}

	// Freeze model at current configuration
	command_torques_.setZero();

	sim->setJointVelocities(kRobotName, dq_des_);
	try {
		readRedisValues();
	} catch (std::exception& e) {
		std::cout << e.what() << std::endl;
	}
	updateModel();
}

/**
 * CSaiEnv::reset()
 * ----------------
 * Reset for OpenAI gym.
 */
void CSaiEnv::reset(uint8_t *observation) {
	reset();

	gl_buffer_ = observation;
	syncGraphics();
};

/**
 * CSaiEnv::step()
 * ---------------
 * Step dynamics for OpenAI gym.
 */
bool CSaiEnv::step(const double *action, uint8_t *observation, double& reward, double *info) {
	for (int i = 0; i < kControlFreq / kEnvironmentFreq; i++) {
		++controller_counter_;
		action_ << action[0], action[1], 0;
		computeOperationalSpaceControlTorques();
		writeRedisValues();
		try {
			readRedisValues();
		} catch (std::exception& e) {
			std::cout << e.what() << std::endl;
		}
		updateModel();
	}

	gl_buffer_ = observation;
	// Tell graphics to update
	mutex_graphics_.lock();
	update_graphics_ = true;
	mutex_graphics_.unlock();
	cv_.notify_all();

	// Compute reward
	if (((x_ - kCenter).array().abs() > kCenterDistance - kWallTolerance).any()) {
		reward = -1;
	} else if ((x_ - (kCenter + Eigen::Vector3d(-kCenterDistance, kCenterDistance, 0))).norm() < kCornerDistance) {
		reward = 1;
	} else {
		reward = 0;
	}

	// Wait for graphics to finish
	std::unique_lock<std::mutex> lock_graphics(mutex_graphics_);
	cv_.wait(lock_graphics, [this]{
		return !update_graphics_;
	});

	// Finish episode after 10s
	bool done = false;//controller_counter_ / kControlFreq >= kTimeEpisode;// ||
				// ((x_ - kCenter).array().abs() > kCenterDistance).any();

	// Return debug info
	if (info != nullptr) {
		for (int i = 0; i < 3; i++) info[  i] = x_(i);
		for (int i = 0; i < 3; i++) info[3+i] = dx_(i);
	}

	return done;
};

/**
 * CSaiEnv::readRedisValues()
 * --------------------------
 * Retrieve all read keys from Redis.
 */
void CSaiEnv::readRedisValues() {
	// Read from Redis current sensor values
	sim->getJointPositions(kRobotName, robot_->_q);
	sim->getJointVelocities(kRobotName, robot_->_dq);
}

/**
 * CSaiEnv::writeRedisValues()
 * ---------------------------
 * Send all write keys to Redis.
 */
void CSaiEnv::writeRedisValues() {
	sim->setJointTorques(kRobotName, command_torques_);
	for (int i = 0; i < kSimulationFreq / kControlFreq; i++) {
		sim->integrate(1.0 / kSimulationFreq);
	}
	redis_.setEigenMatrix(KEY_JOINT_POSITIONS, robot_->_q);
	redis_.setEigenMatrix(KEY_EE_POS, x_);
	redis_.setEigenMatrix(KEY_EE_POS_DES, action_);
}

/**
 * CSaiEnv::updateModel()
 * ----------------------
 * Update the robot model and all the relevant member variables.
 */
void CSaiEnv::updateModel() {
	std::lock_guard<std::mutex> lock(mutex_robot_);

	// Update the model
	robot_->updateModel();

	// Forward kinematics
	x_ = robot_->position("link6", Eigen::Vector3d::Zero());
	dx_ = robot_->linearVelocity("link6", Eigen::Vector3d::Zero());
	w_ = robot_->angularVelocity("link6");
	R_ee_to_base_ = robot_->rotation("link6");

	// Jacobians
	J_ = robot_->J("link6", Eigen::Vector3d::Zero());
	N_ = robot_->nullspaceMatrix(J_);

	// Dynamics
	Lambda_ = robot_->taskInertiaMatrixWithPseudoInv(J_);
	g_ = robot_->gravityVector();
}

/**
 * CSaiEnv::computeJointSpaceControlTorques()
 * ------------------------------------------
 * Controller to initialize robot to desired joint position.
 */
CSaiEnv::ControllerStatus CSaiEnv::computeJointSpaceControlTorques() {
	// Finish if the robot has converged to q_initial
	Eigen::VectorXd q_err = robot_->_q - q_des_;
	Eigen::VectorXd dq_err = robot_->_dq - dq_des_;
	if (q_err.norm() < kToleranceInitQ && dq_err.norm() < kToleranceInitDq) {
		return FINISHED;
	}

	// Compute torques
	Eigen::VectorXd ddq = -kp_joint_ * q_err - kv_joint_ * dq_err;
	command_torques_ = robot_->_M * ddq + g_;
	return RUNNING;
}

/**
 * CSaiEnv::computeOperationalSpaceControlTorques()
 * ------------------------------------------------
 * Controller to move end effector to desired position.
 */
CSaiEnv::ControllerStatus CSaiEnv::computeOperationalSpaceControlTorques() {
	// PD position control with velocity saturation
	Eigen::Vector3d x_err = x_ - x_des_;

	// Push back x if hit wall
	if (x_(0) < kCenter(0) - kCenterDistance) {
		x_err(0) = x_(0) - kCenter(0);
	} else if (x_(0) > kCenter(0) + kCenterDistance) {
		x_err(0) = x_(0) - kCenter(0);
	} else {
		x_err(0) = kp_action_ / kp_pos_ * (action_(0));
	}

	// Push back y if hit wall
	if (x_(1) < kCenter(1) - kCenterDistance) {
		x_err(1) = x_(1) - kCenter(1);
	} else if (x_(1) > kCenter(1) + kCenterDistance) {
		x_err(1) = x_(1) - kCenter(1);
	} else {
		x_err(1) = kp_action_ / kp_pos_ * (action_(1));
	}

	Eigen::Vector3d dx_err = dx_ - dx_des_;
	Eigen::Vector3d ddx = -kp_pos_ * x_err - kv_pos_ * dx_err;
	// dx_des_ = -(kp_pos_ / kv_pos_) * x_err;
	// double v = kMaxVelocity / dx_des_.norm();
	// if (v > 1) v = 1;
	// Eigen::Vector3d dx_err = dx_ - v * dx_des_;
	// Eigen::Vector3d ddx = -kv_pos_ * dx_err;

	// Orientation
	Eigen::Vector3d dPhi = robot_->orientationError(R_des_, R_ee_to_base_);
	Eigen::Vector3d dw = -kp_ori_ * dPhi - kv_ori_ * w_;

	// Nullspace posture control and damping
	Eigen::VectorXd q_err = robot_->_q - q_des_;
	Eigen::VectorXd dq_err = robot_->_dq - dq_des_;
	Eigen::VectorXd ddq = -kp_joint_ * q_err - kv_joint_ * dq_err;

	// Control torques
	Eigen::VectorXd ddx_dw(6);
	ddx_dw << ddx, dw;
	Eigen::VectorXd F = Lambda_ * ddx_dw;
	Eigen::VectorXd F_posture = robot_->_M * ddq;
	command_torques_ = J_.transpose() * F + N_.transpose() * F_posture + g_;

	return RUNNING;
}

CSaiEnv::ControllerStatus CSaiEnv::computeSimpleOperationalSpaceControlTorques() {
	// PD position control with velocity saturation
	Eigen::Vector3d x_err = x_ - x_des_;
	Eigen::Vector3d dx_err = dx_ - dx_des_;
	Eigen::Vector3d ddx = -kp_pos_ * x_err - kv_pos_ * dx_err;

	// Orientation
	Eigen::Vector3d dPhi = robot_->orientationError(R_des_, R_ee_to_base_);
	Eigen::Vector3d dw = -kp_ori_ * dPhi - kv_ori_ * w_;

	// Nullspace posture control and damping
	Eigen::VectorXd q_err = robot_->_q - q_des_;
	Eigen::VectorXd dq_err = robot_->_dq - dq_des_;
	Eigen::VectorXd ddq = -kp_joint_ * q_err - kv_joint_ * dq_err;

	// Control torques
	Eigen::VectorXd ddx_dw(6);
	ddx_dw << ddx, dw;
	Eigen::VectorXd F = Lambda_ * ddx_dw;
	Eigen::VectorXd F_posture = robot_->_M * ddq;
	command_torques_ = J_.transpose() * F + N_.transpose() * F_posture + g_;

	return RUNNING;
}

/**
 * public CSaiEnv::initialize()
 * ----------------------------
 * Initialize timer and Redis client
 */
void CSaiEnv::initialize() {
	// Create a loop timer
	timer_.setLoopFrequency(kControlFreq);   // 1 KHz
	// timer.setThreadHighPriority();  // make timing more accurate. requires running executable as sudo.
	timer_.setCtrlCHandler(stop);    // exit while loop on ctrl-c
	timer_.initializeTimer(kInitializationPause); // 1 ms pause before starting loop

	// Start redis client
	// Make sure redis-server is running at localhost with default port 6379
	redis_.connect(kRedisHostname, kRedisPort);
}

/**
 * public CSaiEnv::syncGraphics()
 * -------------------------------
 * Wait for graphics to finish update.
 */
void CSaiEnv::syncGraphics() {
	// Tell graphics to update
	mutex_graphics_.lock();
	update_graphics_ = true;
	mutex_graphics_.unlock();
	cv_.notify_all();

	// Wait for graphics to finish
	std::unique_lock<std::mutex> lock_graphics(mutex_graphics_);
	cv_.wait(lock_graphics, [this]{
		return !update_graphics_;
	});
}

/**
 * public CSaiEnv::graphicsMain()
 * ------------------------------
 * Graphics thread.
 */
void CSaiEnv::graphicsMain(std::shared_ptr<Graphics::GraphicsInterface> graphics) {
	Graphics::ChaiGraphics *chai = dynamic_cast<Graphics::ChaiGraphics *>(graphics->_graphics_internal);

	// Start visualization
    glfwSetErrorCallback(glfwError);
    glfwInit();

	// Create window
    const GLFWvidmode* mode = glfwGetVideoMode(glfwGetPrimaryMonitor());
    glfwWindowHint(GLFW_VISIBLE, 0);
    GLFWwindow *window = glfwCreateWindow(kWindowWidth, kWindowHeight, "SAI2 OpenAI Gym Environment", nullptr, nullptr);
	glfwSetWindowPos(window, 0.5 * (mode->height - kWindowHeight), 0.5 * (mode->height - kWindowHeight));
	glfwShowWindow(window);
    glfwMakeContextCurrent(window);
	glfwSwapInterval(1);
	chai->setCameraPose(kCameraName, kCameraPos, kCameraVertical, kCameraLookat);

	while (g_runloop && !glfwWindowShouldClose(window)) {
		// Wait for update notification
		std::unique_lock<std::mutex> lock_graphics(mutex_graphics_);
		cv_.wait(lock_graphics, [this]{
			return update_graphics_ || !g_runloop;
		});
		if (!g_runloop) break;

		// Update robot visualization
		mutex_robot_.lock();
		chai->updateGraphics(kRobotName, robot_.get());
		mutex_robot_.unlock();

		// Render graphics
		chai->render(kCameraName, kWindowWidth, kWindowHeight);
		glfwSwapBuffers(window);
		glFinish();

		// Take screenshot
		if (gl_buffer_ == nullptr) {
			gl_buffer_ = new uint8_t[3 * kWindowWidth * kWindowHeight];
		}
		glReadPixels(0, 0, kWindowWidth, kWindowHeight, GL_RGB, GL_UNSIGNED_BYTE, gl_buffer_);
		if (glGetError() != GL_NO_ERROR) g_runloop = false;

		// Notify threads waiting on graphics
		update_graphics_ = false;
		lock_graphics.unlock();
		cv_.notify_all();
	}
	g_runloop = false;

	glfwDestroyWindow(window);
	glfwTerminate();
}


/**
 * public CSaiEnv::runLoop()
 * -------------------------
 * CSaiEnv state machine
 */
void CSaiEnv::runLoop() {

	while (g_runloop) {
		// Wait for next scheduled loop (controller must run at precise rate)
		// timer_.waitForNextLoop();
		++controller_counter_;

		// Get latest sensor values from Redis and update robot model
		try {
			readRedisValues();
		} catch (std::exception& e) {
			if (controller_state_ != REDIS_SYNCHRONIZATION) {
				std::cout << e.what() << " Aborting..." << std::endl;
				break;
			}
			std::cout << e.what() << " Waiting..." << std::endl;
			std::this_thread::sleep_for(std::chrono::seconds(1));
			continue;
		}
		updateModel();

		switch (controller_state_) {
			// Wait until valid sensor values have been published to Redis
			case REDIS_SYNCHRONIZATION:
				if (isnan(robot_->_q)) continue;
				std::cout << "Redis synchronized. Switching to joint space controller." << std::endl;
				controller_state_ = JOINT_SPACE_INITIALIZATION;
				break;

			// Initialize robot to default joint configuration
			case JOINT_SPACE_INITIALIZATION:
				if (computeJointSpaceControlTorques() == FINISHED) {
					std::cout << "Joint position initialized. Switching to operational space controller." << std::endl;
					controller_state_ = CSaiEnv::OP_SPACE_POSITION_CONTROL;
				};
				break;

			// Control end effector to desired position
			case OP_SPACE_POSITION_CONTROL:
				computeOperationalSpaceControlTorques();
				break;

			// Invalid state. Zero torques and exit program.
			default:
				std::cout << "Invalid controller state. Stopping controller." << std::endl;
				g_runloop = false;
				command_torques_.setZero();
				break;
		}

		// Check command torques before sending them
		if (isnan(command_torques_)) {
			std::cout << "NaN command torques. Sending zero torques to robot." << std::endl;
			command_torques_.setZero();
		}

		// Send command torques
		writeRedisValues();

		if (controller_counter_ % (kControlFreq / kEnvironmentFreq) == 0) {
			mutex_graphics_.lock();
			update_graphics_ = true;
			mutex_graphics_.unlock();
			cv_.notify_all();
		}
	}

	// Zero out torques before quitting
	command_torques_.setZero();
	// redis_.setEigenMatrix(KEY_COMMAND_TORQUES, command_torques_);

	thread_graphics.join();
}

int main(int argc, char** argv) {

	// Parse command line
	if (argc != 4) {
		std::cout << "Usage: demo_app <path-to-world.urdf> <path-to-robot.urdf> <robot-name>" << std::endl;
		exit(0);
	}
	// Argument 0: executable name
	// Argument 1: <path-to-world.urdf>
	std::string world_file(argv[1]);
	// Argument 2: <path-to-robot.urdf>
	std::string robot_file(argv[2]);
	// Argument 3: <robot_-name>
	std::string robot_name(argv[3]);

	// Set up signal handler
	signal(SIGABRT, &stop);
	signal(SIGTERM, &stop);
	signal(SIGINT, &stop);

	// Load robot
	std::cout << "Loading robot: " << robot_file << std::endl;
	auto robot = std::make_shared<Model::ModelInterface>(robot_file, Model::rbdl, Model::urdf, false);
	auto sim = std::make_shared<Simulation::SimulationInterface>(world_file, Simulation::sai2simulation, Simulation::urdf, false);
	auto graphics = std::make_shared<Graphics::GraphicsInterface>(world_file, Graphics::chai, Graphics::urdf, true);

	// Start controller app
	std::cout << "Initializing app with " << robot_name << std::endl;
	CSaiEnv app(std::move(robot), std::move(sim), std::move(graphics), robot_name);
	app.initialize();
	std::cout << "App initialized. Waiting for Redis synchronization." << std::endl;
	app.runLoop();

	return 0;
}


<file_sep>from gym_sai2.envs.sai2_env import SaiEnv
<file_sep>set -e

# ---------------------------------------
# Install precompiled 3rd party libraries
# ---------------------------------------

if [[ "$OSTYPE" == "linux-gnu" ]]; then
	sudo apt-get install curl cmake libeigen3-dev libtinyxml2-dev libjsoncpp-dev libhiredis-dev libglfw3-dev xorg-dev freeglut3-dev libasound2-dev libusb-1.0-0-dev redis-server
	# Install gcc 5 for Ubuntu 14.04:
	# sudo add-apt-repository ppa:ubuntu-toolchain-r/test
	# sudo apt-get update
	# sudo apt-get install gcc-5 g++-5
	# sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-5 60 --slave /usr/bin/g++ g++ /usr/bin/g++-5
elif [[ "$OSTYPE" == "darwin"* ]]; then
	brew install cmake eigen redis hiredis tinyxml2 jsoncpp glfw3
fi

mkdir -p lib
cd lib

./download.sh

cd ..
<file_sep>"""
Utility functions for processing data.
"""
import os
import h5py
import numpy as np
import math
import glob
from time import strftime, gmtime

def list_files():
    filepath = os.path.join("data")
    files = glob.glob(os.path.join(filepath, "*.hdf5"))
    files = sorted(files, key=lambda f: os.stat(f).st_mtime)
    return files

def get_filename(idx):
    filename = list_files()[idx]
    print("Loading: ", filename)
    return filename

def get_mean_observation(filename):
    f = h5py.File(filename, "r")

    # Calculate mean observation
    if not "mean_observation" in f:
        dim_o = np.array(f["episodes/00000/observations"].shape)[1:]
        
        # Initialize mean
        D = 100
        D_actual = 0
        mean_observation = np.zeros(dim_o)

        # Iterate over episodes
        episodes = f["episodes"]
        for i in episodes:
            try:
                observations = episodes[i]["observations"][()]
            except:
                # Corrupt data
                continue

            observation = (observations.astype(np.float64) / 255).mean(axis=0)
            mean_observation += observation / D
            D_actual += 1
            # if D_actual >= D:
            #     break

        mean_observation *= D / D_actual

        # Write mean observation to file
        f.close()
        with h5py.File(filename, "a") as f:
            dset = f.create_dataset("mean_observation", mean_observation.shape, dtype=mean_observation.dtype)
            dset[...] = mean_observation

    else:
        mean_observation = f["mean_observation"][()]
        f.close()
    return mean_observation

def get_observation_dim(filename=None):
    if filename is None:
        filename = get_filename(-1)

    mean_observation = get_mean_observation(filename)
    return mean_observation.shape

def raw_data(size_batch=100, filename=None):
    if filename is None:
        filename = get_filename(-1)

    f = h5py.File(filename, "r")

    episodes = f["episodes"]

    o = []
    T = 0
    for i in range(len(episodes)):
        grp = episodes["{0:05d}".format(i)]
        observations = grp["observations"][()].astype(np.float32)

        # Append data to lists
        o.append(observations)
        T += observations.shape[0]

        if T >= size_batch:
            # Gather batch
            o = np.concatenate(o, axis=0)
            return o[:size_batch]


def batch_data(data=None, size_batch=100, extra=False, filename=None, dataset="all", flatten=True):
    if data is not None:
        observations = data[1]
        try:
            mean_observation = np.load("../resources/mean_observation.npy")
        except:
            print("Could not load ../resources/mean_observation.npy")
            num_observations = sum(o.shape[0] for o in observations)
            mean_observation = sum(np.mean(o.astype(np.float64), axis=0) * (o.shape[0] / num_observations / 255) for o in observations)
            mean_observation = mean_observation[np.newaxis,...].astype(np.float32)

        for a, o, r, x, dx, s_hat in zip(*data):
            # Preprocess observation
            o = o.astype(np.float32) / 255 - mean_observation
            if flatten:
                o = o.reshape(o.shape[0], -1)
                assert np.all(np.any(np.abs(o) > 0.01, axis=1))

            if extra:
                x = x[...,:2] - np.array([0, -0.45])[np.newaxis,:]
                dx = dx[...,:2]
                yield (o, a, r, x, dx, s_hat)
            else:
                yield (o, a, r)

        return

    if filename is None:
        filename = get_filename(-1)

    mean_observation = get_mean_observation(filename)
    f = h5py.File(filename, "r")

    episodes = f["episodes"]
    if dataset == "train":
        D = (0, math.floor(0.8 * len(episodes)))
        size_batch = min(1000 * D[1], size_batch)
    elif dataset == "test":
        D = (math.floor(0.8 * len(episodes)), len(episodes))
        size_batch = min(1000 * (D[1] - D[0]), size_batch)
    else:
        D = (0, len(episodes))
        size_batch = min(1000 * D[1], size_batch)
    print(dataset, D, len(episodes))

    while True:

        if D[0] == 0:
            o_0 = f["initial_observation"][()].astype(np.float32)
            if extra:
                x_0 = np.zeros(2, dtype=np.float32)
        else:
            try:
                o_0 = episodes["{0:05d}/observations".format(D[0]-1)][-1,...].astype(np.float32)
                if extra:
                    x_0 = episodes["{0:05d}/xs".format(D[0]-1)][-1,...].astype(np.float32)
                    x_0 -= np.array([0, -0.45], dtype=np.float32)[np.newaxis,:]
            except:
                o_0 = f["initial_observation"][()].astype(np.float32)
                if extra:
                    x_0 = np.zeros(2, dtype=np.float32)
        o_0 = o_0 / 255 - mean_observation
        if flatten:
            o_0 = o_0.reshape(-1)
        o_0 = o_0[np.newaxis,...]
        o  = [o_0]
        a  = []
        r  = []
        if extra:
            x_0 = x_0[np.newaxis,:]
            x  = [x_0]
            dx = []
        T = 0

        for i in range(D[0], D[1]):
            try:
                grp = episodes["{0:05d}".format(i)]
                actions      = grp["actions"][()].astype(np.float32)
                observations = grp["observations"][()].astype(np.float32)
                rewards      = grp["rewards"][()].astype(np.float32)
                if extra:
                    xs  = grp["xs"][()].astype(np.float32)
                    dxs = grp["dxs"][()].astype(np.float32)
            except:
                continue

            # Flatten observations and center
            observations = observations / 255 - mean_observation[np.newaxis,...]
            if flatten:
                observations = observations.reshape(observations.shape[0],-1)
                assert np.all(np.any(np.abs(observations) > 0.01, axis=1))

            # Append data to lists
            o.append(observations)
            a.append(actions)
            r.append(rewards)
            if extra:
                x.append(xs[...,:2] - np.array([0, -0.45], dtype=np.float32)[np.newaxis,:])
                dx.append(dxs[...,:2])
            T += observations.shape[0]

            if T >= size_batch:
                # Gather batch
                o  = np.concatenate(o, axis=0)
                a  = np.concatenate(a, axis=0)
                r  = np.concatenate(r, axis=0)
                if extra:
                    x  = np.concatenate(x, axis=0)
                    dx = np.concatenate(dx, axis=0)

                # Yield batch
                if dataset == "test":
                    while True:
                        if extra:
                            yield (o, a, r, x, dx)
                        else:
                            yield (o, a, r)
                else:
                    if extra:
                        yield (o, a, r, x, dx)
                    else:
                        yield (o, a, r)
                o_0 = o[-1,np.newaxis,...]
                o  = [o_0]
                a  = []
                r  = []
                if extra:
                    x_0 = x[-1,np.newaxis,:]
                    x  = [x_0]
                    dx = []
                T = 0


        # print("Dataset finished: " + dataset)

class DataLogger:

    def __init__(self):
        if not os.path.exists("results"):
            os.makedirs("results")
        self.filename = "results/data-{}.hdf5".format(strftime("%m-%d_%H-%M"), gmtime())
        self.f = h5py.File(self.filename, "w")

        self.actions_history = []
        self.observations_history = []
        self.rewards_history = []
        self.xs_history = []
        self.dxs_history = []
        self.learned_states_history = []

    def log_initial_observation(self, initial_observation):
        dset = self.f.create_dataset("initial_observation", initial_observation.shape, dtype=initial_observation.dtype)
        dset[...] = initial_observation

    def log(self, i, actions, observations, rewards, xs, dxs, learned_states):
        # Save trajectory to dataset
        grp = self.f.create_group("/episodes/{0:05d}".format(i))
        dset = grp.create_dataset("actions", actions.shape, dtype=actions.dtype)
        dset[...] = actions
        dset = grp.create_dataset("observations", observations.shape, dtype=observations.dtype)
        dset[...] = observations
        dset = grp.create_dataset("rewards", rewards.shape, dtype=rewards.dtype)
        dset[...] = rewards
        dset = grp.create_dataset("xs", xs.shape, dtype=xs.dtype)
        dset[...] = xs
        dset = grp.create_dataset("dxs", dxs.shape, dtype=dxs.dtype)
        dset[...] = dxs
        dset = grp.create_dataset("learned_states", learned_states.shape, dtype=learned_states.dtype)
        dset[...] = learned_states

        # Add to history
        self.actions_history.append(actions)
        self.observations_history.append(observations)
        self.rewards_history.append(rewards)
        self.xs_history.append(xs)
        self.dxs_history.append(dxs)
        self.learned_states_history.append(learned_states)

    def flush(self):

        actions = self.actions_history
        observations = self.observations_history
        rewards = self.rewards_history
        xs = self.xs_history
        dxs = self.dxs_history
        learned_states = self.learned_states_history

        self.actions_history = []
        self.observations_history = []
        self.rewards_history = []
        self.xs_history = []
        self.dxs_history = []
        self.learned_states_history = []

        self.f.flush()

        return actions, observations, rewards, xs, dxs, learned_states

    def __enter__(self):
        return self

    def __exit__(self, etype, value, traceback):
        self.f.close()
<file_sep>"""
Variational autoencoder (experimental).
"""
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from data import *
import os
from time import strftime
from robotic_priors import *

print(list_files())
filename = get_filename(-1)
print(filename)

train_batch = batch_data(size_batch=1000, extra=True, filename=filename, dataset="train", reshape=False)
test_batch  = batch_data(size_batch=5000, extra=True, filename=filename, dataset="test", reshape=False)
x_test, _, _, _, _ = next(test_batch)

dim_image = list(get_observation_dim(filename))
dim_o = 2*28224#np.prod(dim_image)
dim_x = 2
dim_a = 2
dim_h1 = 32
dim_h2 = 8
krn_h1 = 5
krn_h2 = 5
input_dim = np.prod(dim_image)#np.prod(np.array(get_observation_dim(filename)))
hidden_encoder_dim = 400
hidden_decoder_dim = 1024
latent_dim = 2
lam = 0

def weight_variable(shape):
  initial = tf.truncated_normal(shape, stddev=0.001)
  return tf.Variable(initial)

def bias_variable(shape):
  initial = tf.constant(0., shape=shape)
  return tf.Variable(initial)

x = tf.placeholder("float", shape=[None] + dim_image, name="x")
x_flat = tf.reshape(x, [-1, input_dim])
l2_loss = tf.constant(0.0)

conv1 = tf.layers.conv2d(
    inputs=x,
    filters=dim_h1,
    kernel_size=(krn_h1, krn_h1),
    padding="valid",
    strides=(2,2),
    activation=tf.nn.relu,
    name="conv1"
)
# 98 x 73 x 16
pool1 = tf.layers.max_pooling2d(
    inputs=conv1,
    pool_size=(2,2),
    strides=2,
    name="pool1"
)
# 49 x 36 x 16
h1 = tf.reshape(pool1, [-1, dim_o])

W_encoder_input_hidden = weight_variable([dim_o,hidden_encoder_dim])
b_encoder_input_hidden = bias_variable([hidden_encoder_dim])
l2_loss += tf.nn.l2_loss(W_encoder_input_hidden)

# Hidden layer encoder
# hidden_encoder = tf.nn.relu(tf.matmul(x, W_encoder_input_hidden) + b_encoder_input_hidden)
hidden_encoder = tf.nn.relu(tf.matmul(h1, W_encoder_input_hidden) + b_encoder_input_hidden)

W_encoder_hidden_mu = weight_variable([hidden_encoder_dim,latent_dim])
b_encoder_hidden_mu = bias_variable([latent_dim])
l2_loss += tf.nn.l2_loss(W_encoder_hidden_mu)

# Mu encoder
mu_encoder = tf.add(tf.matmul(hidden_encoder, W_encoder_hidden_mu), b_encoder_hidden_mu, name= "mu_encoder")

W_encoder_hidden_logvar = weight_variable([hidden_encoder_dim,latent_dim])
b_encoder_hidden_logvar = bias_variable([latent_dim])
l2_loss += tf.nn.l2_loss(W_encoder_hidden_logvar)

# Sigma encoder
logvar_encoder = tf.matmul(hidden_encoder, W_encoder_hidden_logvar) + b_encoder_hidden_logvar

# Sample epsilon
epsilon = tf.random_normal(tf.shape(logvar_encoder), name='epsilon')

# Sample latent variable
std_encoder = tf.exp(0.5 * logvar_encoder)
z = mu_encoder + tf.multiply(std_encoder, epsilon)

W_decoder_z_hidden = weight_variable([latent_dim,hidden_decoder_dim])
b_decoder_z_hidden = bias_variable([hidden_decoder_dim])
l2_loss += tf.nn.l2_loss(W_decoder_z_hidden)

# Hidden layer decoder
hidden_decoder = tf.nn.relu(tf.matmul(z, W_decoder_z_hidden) + b_decoder_z_hidden)

W_decoder_hidden_reconstruction = weight_variable([hidden_decoder_dim, input_dim])
b_decoder_hidden_reconstruction = bias_variable([input_dim])
l2_loss += tf.nn.l2_loss(W_decoder_hidden_reconstruction)

KLD = -0.5 * tf.reduce_sum(1 + logvar_encoder - tf.pow(mu_encoder, 2) - tf.exp(logvar_encoder), reduction_indices=1)

x_hat = tf.add(tf.matmul(hidden_decoder, W_decoder_hidden_reconstruction), b_decoder_hidden_reconstruction, name="x_hat")
BCE = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=x_hat, labels=x_flat), reduction_indices=1)

loss = tf.reduce_mean(BCE + KLD)

regularized_loss = loss + lam * l2_loss

loss_sum = tf.summary.scalar("lowerbound", loss)
train_step = tf.train.AdamOptimizer(0.001).minimize(regularized_loss)

# add op for merging summary
summary_op = tf.summary.merge_all()

# add Saver ops
saver = tf.train.Saver()

n_steps = int(1e6)
batch_size = 100

datadir = os.path.join("data", "logs", strftime("%m-%d_%H-%M"))
modeldir = os.path.join("data", "models", strftime("%m-%d_%H-%M"))
modelpath = os.path.join(modeldir, "model")
train_writer = tf.summary.FileWriter(os.path.join(datadir, "train"))
test_writer = tf.summary.FileWriter(os.path.join(datadir, "test"))
saver = tf.train.Saver()
min_loss_test = float("inf")

with tf.Session() as sess:
  summary_writer = tf.summary.FileWriter('experiment',
                                          graph=sess.graph)
  if os.path.isfile("save/model.ckpt"):
    print("Restoring saved parameters")
    saver.restore(sess, "save/model.ckpt")
  else:
    print("Initializing parameters")
    sess.run(tf.global_variables_initializer())

  for step in range(1000):
    x_train, _, _, _, _ = next(train_batch)
    _, loss_train, summary_train = sess.run([train_step, loss, summary_op], feed_dict={x: x_train})
    loss_test, summary_test = sess.run([loss, summary_op], feed_dict={x: x_test})

    if step % 100 == 0:
      train_writer.add_summary(summary_train, step)
      test_writer.add_summary(summary_test, step)
    if loss_test < min_loss_test:
      min_loss_test = loss_test
      save_path = saver.save(sess, modelpath, global_step=step)
    print("Step {0} | Loss: {1}, {2}".format(step, loss_train, loss_test))

  train_writer.close()
  test_writer.close()
