Categories
3D Research AI/ML CNNs deep dev envs evolution GANs Gripper Gripper Research Linux Locomotion sexing sim2real simulation The Sentient Table UI Vision

Simulation Vision

We’ve got an egg in the gym environment now, so we need to collect some data for training the robot to go pick up an egg.

I’m going to have it save the rgba, depth and segmentation images to disk for Unet training. I left out the depth image for now. The pictures don’t look useful. But some papers are using the depth, so I might reconsider. Some weed bot paper uses 14-channel images with all sorts of extra domain specific data relevant to plants.

I wrote some code to take pics if the egg was in the viewport, and it took 1000 rgb and segmentation pictures or so. I need to change the colour of the egg for sure, and probably randomize all the textures a bit. But main thing is probably to make the segmentation layers with pixel colours 0,1,2, etc. so that it detects the egg and not so much the link in the foreground.

So sigmoid to softmax and so on. Switching to multi-class also begs the question whether to switch to Pytorch & COCO panoptic segmentation based training. It will have to happen eventually, as I think all of the fastest implementations are currently in Pytorch and COCO based. Keras might work fine for multiclass or multiple binary classification, but it’s sort of the beginning attempt. Something that works. More proof of concept than final implementation. But I think Keras will be good enough for these in-simulation 256×256 images.

Regarding multi-class segmentation, karolzak says “it’s just a matter of changing num_classes argument and you would need to shape your mask in a different way (layer per class??), so for multiclass segmentation you would need a mask of shape (width, height, num_classes)

I’ll keep logging my debugging though, if you’re reading this.

So I ran segmask_linkindex.py to see what it does, and how to get more useful data. The code is not running because the segmentation image actually has an array of arrays. I presume it’s a numpy array. I think it must be the rows and columns. So anyway I added a second layer to the loop, and output the pixel values, and when I ran it in the one mode:

-1
-1
-1
83886081
obUid= 1 linkIndex= 4
83886081
obUid= 1 linkIndex= 4
1
obUid= 1 linkIndex= -1
1
obUid= 1 linkIndex= -1
16777217
obUid= 1 linkIndex= 0
16777217
obUid= 1 linkIndex= 0
-1
-1
-1

And in the other mode

-1
-1
-1
1
obUid= 1 linkIndex= -1
1
obUid= 1 linkIndex= -1
1
obUid= 1 linkIndex= -1
-1
-1
-1

Ok I see. Hmm. Well the important thing is that this code is indeed for extracting the pixel information. I think it’s going to be best for the segmentation to use the simpler segmentation mask that doesn’t track the link info. Ok so I used that code from the guy’s thesis project, and that was interpolating the numbers. When I look at the unique elements of the mask without interpolation, I’ve got…

[  0   2 255]
[  0   2 255]
[  0   2 255]
[  0   2 255]
[  0   2 255]
[  0   1   2 255]
[  0   1   2 255]
[  0   2 255]
[  0   2 255]

Ok, so I think:

255 is the sky
0 is the plane
2 is the robotable
1 is the egg

So yeah, I was just confused because the segmentation masks were all black and white. But if you look closely with a pixel picker tool, the pixel values are (0,0,0), (1,1,1), (2,2,2), (255,255,255), so I just couldn’t see it.

The interpolation kinda helps, to be honest.

As per OpenAI’s domain randomization helping with Sim2Real, we want to randomize some textures and some other things like that. I also want to throw in some random chickens. Maybe some cats and dogs. I’m afraid of transfer learning, at this stage, because a lot of it has to do with changing the structure of the final layer of the neural network, and that might be tough. Let’s just do chickens and eggs.

An excerpt from OpenAI:

Costs

Both techniques increase the computational requirements: dynamics randomization slows training down by a factor of 3x, while learning from images rather than states is about 5-10x slower.

Ok that’s a bit more complex than I was thinking. I want to randomize textures and colours, first

I’ve downloaded and unzipped the ‘Describable Textures Dataset’

And ok it’s loading a random texture for the plane

and random colour for the egg and chicken

Ok, next thing is the Simulation CNN.

Interpolation doesn’t work though, for this, cause it interpolates from what’s available in the image:

[  0  85 170 255]
[  0  63 127 191 255]
[  0  63 127 191 255]

I kind of need the basic UID segmentation.

[  0   1   2   3 255]

Ok, pity about the mask colours, but anyway.

Let’s train the UNet on the new dataset.

We’ll need to make karolzak’s changes.

I’ve saved 2000+ rgb.jpg and seg.png files and we’ve got [0,1,2,3,255] [plane, egg, robot, chicken, sky]

So num_classes=5

And

“for multiclass segmentation you would need a mask of shape (width, height, num_classes) “

What is y.shape?

(2001, 256, 256, 1)

which is 2001 files, of 256 x 256 pixels, and one class. So if I change that to 5…? ValueError: cannot reshape array of size 131137536 into shape (2001,256,256,5)

Um… Ok I need to do more research. Brb.

So the keras_unet library is set up to input binary masks per class, and output binary masks per class.

I would rather use the ‘integer’ class output, and have it output a single array, with the class id per pixel. Similar to this question. In preparation for karolzak probably not knowing how to do this with his library, I’ve asked on stackoverflow for an elegant way to make the binary masks from a multi-class mask, in the meantime.

I coded it up using the library author’s suggested method, as he pointed out that the gains of the integer encoding method are minimal. I’ll check it out another time. I think it might still make sense for certain cases.

Ok that’s pretty awesome. We have 4 masks. Human, chicken, egg, robot. I left out plane and sky for now. That was just 2000 images of training, and I have 20000. I trained on another 2000 images, and it’s down to 0.008 validation loss, which is good enough!

So now I want to load the CNN model in the locomotion code, and feed it the images from the camera, and then have a reward function related to maximizing the egg pixels.

I also need to look at the pybullet-planning project and see what it consists of, as I imagine they’ve made some progress on the next steps. “built-in implementations of standard motion planners, including PRM, RRT, biRRT, A* etc.” – I haven’t even come across these acronyms yet! Ok, they are motion planning. Solvers of some sort. Hmm.

Categories
3D 3D prototypes 3D Research control dev envs Gripper gripper prototypes Gripper Research Locomotion robots Vision

Gripper simulation

I’ve been scouring for existing code to help with developing the gripper in simulation. I was looking for a way to implement ‘eye-in-hand’ visual servoing, and came across a good resource, created for a masters thesis, which shows a ‘robot vision’ window, and he compares depth sensing algorithms. My approach was going to be, essentially, segmentation, in order to detect and localise chickens and eggs, in the field of vision, and then just try get their shape into an X-Y coordinate position, and over a certain size, to initiate interaction.

This one uses an SDF model of a KUKA industrial 6 DOF robot with a two finger gripper, but that has specific rotational movement, that seems maybe different from a simpler robot arm. So it’s maybe a bit overkill, and I might just want to see his camera code.

Miranda’s gripper prototype isn’t a $50k KUKA industrial robot arm. It’s just v.0.1 and got an 11kg/cm MG945, some 5kg/cm MG5010s, and an 1.3kg/cm SG90, and a sucker contraption I found on DFRobot, that can suck eggs.

So, regarding the simulation,this will be on top of the robot, as its head.

So we need an URDF file. Or an SDF file. There’s a couple ways to go with this.

The other resource I’ve found that looks like just what I need, is ur5pybullet

Regarding the ‘visual servoing’, the state of the art appears to be QT-Opt, perhaps. Or maybe RCAN, built on top of it. But we’re not there just yet. Another project specifically uses pybullet. Some extra notes here, from Sergey Levine, and co., associated with most of these projects.

Another good one is Retina-GAN, where they convert both simulation and reality into a canonical format. I’ve also come across Dex-Net before, from UCB.

Everything is very complicated though.

I’ve managed to make an URDF that looks good enough to start with, though. I’ll put everything in a github. We want to put two servos on the ‘head’ for animatronic emotional aesthetics, but there’s a sucker contraption there for the egg, so I think this is good enough for simulation, for now, anyway. I just need to put a camera on its head, put some eggs in the scene, and maybe reward stable contact with the tip. Of course it’s going to be a lot of work.

We also want to add extra leg parts, but I don’t want to use 4 more motors on it.

So I’m playing around with some aluminium and timing belts and pulleys to get 8 leg parts on 4 motors. Something like this, with springs if we can find some.

So, simulator camera vision. I can enable the GUI. Turns out I just need to press ‘g’ to toggle.

self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_RENDERING, 1)
self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_GUI, 1)
self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 1)
self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 1)
self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_RGB_BUFFER_PREVIEW, 1)

I’ve added the gripper, and now I’m printing out the _control_observation, because I need to work out what is in an observation.

self.GetTrueMotorAngles()
0.18136442583543283, 0.4339093246887722, -0.25269494256467184, 0.32002873424829736, -0.6635045784503064, 1.5700002984158676, -1.5700000606174402, -0.2723645141027962,

self.GetTrueMotorVelocities()
0.451696256678765, 0.48232988947216504, -4.0981980703534395, 0.4652986924553241, 0.3592921211587608, -6.978131098967118e-06, 1.5237597481713495e-06, -10.810712328063294,

self.GetTrueMotorTorques()
-3.5000000000000004, -3.5000000000000004, 3.5000000000000004, -3.5000000000000004, 3.5000000000000004, -3.5000000000000004, 3.5000000000000004, 3.5000000000000004, 

self.GetTrueBaseOrientation()
-0.008942336195953221, -0.015395612988274186, 0.00639837318132646, 0.9998210192552996, 

self.GetTrueBaseRollPitchYawRate()
-0.01937158793669886, -0.05133982438770338, 0.001050170752804882]

Ok so I need the link state of the end effector (8th link), to get its position and orientation.

    state = self._pybullet_client.getLinkState(self.quadruped, self._end_effector_index)
    pos = state[0]
    orn = state[1]

    print(pos)
    print(orn)

(0.8863188372297804, -0.4008813832608453, 3.1189486984341848)

(0.9217446940545668, 0.3504950513334899, -0.059006227834041206, -0.1551070696318658)

Since the orientation is 4 dimensions, it’s a quaternion,

  def gripper_camera(self, state):
    pos = state[0]
    ori = state[1]


    rot_matrix = self._pybullet_client.getMatrixFromQuaternion(ori)

    rot_matrix = np.array(rot_matrix).reshape(3, 3)
    
# Initial vectors
    init_camera_vector = (1, 0, 0) # z-axis
    init_up_vector = (0, 1, 0) # y-axis
    
# Rotated vectors
    camera_vector = rot_matrix.dot(init_camera_vector)
    up_vector = rot_matrix.dot(init_up_vector)

    self.view_matrix_gripper = self._pybullet_client.computeViewMatrix(pos, pos + 0.1 * camera_vector, up_vector)

    img = self._pybullet_client.getCameraImage(256, 256, self.view_matrix_gripper, self.projectionMatrix, shadow=0, flags = self._pybullet_client.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX, renderer=self._pybullet_client.ER_BULLET_HARDWARE_OPENGL)

Ok I’ve got the visuals now, but I shouldn’t be seeing that shadow

The camera is like 90 degrees off maybe. Could be an issue with the camera setup, or maybe the URDF setup? Ok…

Changing the initial camera vector fixed the view somewhat:

    init_camera_vector = (0, 0, 1) # x-axis

Except that we’re looking backwards now.

init_camera_vector = (0, 0, -1) # x-axis

Ok well it’s correct now, but heh, hmm. Might need to translate the camera just a bit higher.

I found a cool free chicken obj file with Creative commons usage. And an egg.

Heh need to resize obj files. Collision physics is fun.

Ok I worked out how to move the camera a bit higher.

    pos = list(pos) 
    pos[2] += 0.3
    pos = tuple(pos)

Alright! Getting somewhere.

So, next, I add resized eggs and some chickens for good measure, to the scene.

Then we need to train it to stick its shnoz on the eggs.

Ok… gonna have to train this sucker now.

First, the table is falling from the sky, so I might need to stabilize it first. I also need to randomize the egg location a bit.

And I want to minimize the distance between the gripper attachment and the egg.

The smart way is probably to have it walk until some condition and then grasp, but in the spirit of letting the robot learn things by itself, I will probably ignore heuristics. If I do decide to use heuristics, it will probably be a finite state machine with ‘walking’ mode and ‘gripping’ mode. But we’ll come back to this when it’s necessary. Most of the time there won’t be any eggs in sight. So it will just need to walk around until it is sure there is an egg somewhere in sight.

Ok I’ve added a random egg to the scene

self.rng = default_rng()
egg_position = np.r_[self.rng.uniform(-20, 20, 2), 0.1]
egg_orientation = transformations.random_quaternion(self.rng.random(3))
self._egg_mesh = self._pybullet_client.loadURDF("%s/egg.urdf" % self._urdf_root, egg_position, egg_orientation, globalScaling=0.1)

And the end effector’s position should be something like the original camera position before we moved it up a bit, plus length of the end effector in the URDF (0.618). I ended up doing this:

    pos = [pos[0] + 0.5*camera_vector[0], 
           pos[1] + 0.5*camera_vector[1], 
           pos[2] + 0.5*camera_vector[2]]
    pos = tuple(pos) 

And it’s closer to the tip now. But yeah. I will start a new post, Simulation Vision.

Categories
AI/ML envs Vision

NVIDIA Tests

Ok here we go. I already installed gstreamer. Was still getting some h264 encoding plugin missing error, so I needed to add this:

apt-get install gstreamer1.0-libav

Then on my Ubuntu laptop (192.168.0.103), I run:

gst-launch-1.0 -v udpsrc port=1234 caps = "application/x-rtp, media=(string)video, clock-rate=(int)90000, encoding-name=(string)H264, payload=(int)96" ! rtph264depay ! decodebin ! videoconvert ! autovideosink

And then to get webcam streaming from the Jetson,

video-viewer /dev/video0 rtp://192.168.0.103:1234

and similarly,

detectnet /dev/video0 rtp://192.168.0.103:1234

and

segnet /dev/video0 rtp://192.168.0.103:1234

It is impractical to run VNC because of the tiny resolution, and ssh -X tunnelling because that requires the host to have whatever drivers are used on the jetson. GStreamer is working well though.

Cool.

I’m trying to run a python program on the feed. Ended up finding same issue elsewhere. RTP output not working. Bumped the thread.

Someone worked it out. Needed a do/while loop instead of a for loop.

It’s been a couple of months since I was here, and it’s now time to load up this gstreamer code again, and see if I can stream a webcam feed from the robot, and we evaluate inference on the trained CNN, and we colour in the different classes in pretty colours, and

So, TODO:

  1. Get panoptic segmentation colour mapping combining the output layers of the CNN. So something like this. We get a live webcam feed, and we run the frames through the semantic segmentation CNN, and combine the binary masks into a multiclass mask.
  2. Then gstream these rainbow multiclass masked images to a listening gstreamer video-viewer program

So, practically though, I need to set up the Jetson again, and get it compiling the trained h5 file, and using webcam stills as input. Then gstream the result. Ok.

And fix power cable for Jetson step 1. Then try set up webcam gstreaming.

Ok, it’s plugged in, nmap found it. Let’s log in… Ok, run gstreamer, and then i’ve got a folder, jetson-inference, and I’m volume mapping it.

./docker/run.sh --volume /home/chicken/jetson-inference:/jetson-inference
import jetson.inference
import jetson.utils

net = jetson.inference.detectNet("ssd-mobilenet-v2", threshold=0.5)

camera = jetson.utils.videoSource("/dev/video0")
display = jetson.utils.videoOutput("rtp://192.168.101.127:1234","--headless") # 'my_video.mp4' for file

while True:
    img = camera.Capture()
    detections = net.Detect(img)
    display.Render(img)
    display.SetStatus("Object Detection | Network {:.0f} FPS".format(net.GetNetworkFPS()))
    if not camera.IsStreaming() or not display.IsStreaming():
        break

I am in this dusty jetson-inference docker.

pip3 list

appdirs (1.4.4)
boto3 (1.16.58)
botocore (1.19.58)
Cython (0.29.21)
dataclasses (0.8)
decorator (4.4.2)
future (0.18.2)
jmespath (0.10.0)
Mako (1.1.3)
MarkupSafe (1.1.1)
numpy (1.19.4)
pandas (1.1.5)
Pillow (8.0.1)
pip (9.0.1)
pycuda (2020.1)
python-dateutil (2.8.1)
pytools (2020.4.3)
pytz (2020.5)
s3transfer (0.3.4)
setuptools (51.0.0)
six (1.15.0)
torch (1.6.0)
torchaudio (0.6.0a0+f17ae39)
torchvision (0.7.0a0+78ed10c)
urllib3 (1.26.2)
wheel (0.36.1)

and my nice docker with TF2 and everything installed already says:

root@jetson:/dmc/jetson-inference/build/aarch64/bin# ./my-detection.py
Traceback (most recent call last):
File "./my-detection.py", line 24, in
import jetson.inference
ModuleNotFoundError: No module named 'jetson'

Ok let’s try install jetson from source. First, tried the ‘Quick Reference’ instructions… Errored at

/dmc/jetson-inference/c/depthNet.h(190): error: identifier "COLORMAP_VIRIDIS_INVERTED" is undefined

/dmc/jetson-inference/c/depthNet.h(180): error: identifier "COLORMAP_VIRIDIS_INVERTED" is undefined

Next, ran a command mentioned lower down,

git submodule update --init

Now make -j$(nproc) gets to

/usr/bin/ld: cannot find -lnvcaffe_parser

And this reply suggests using sed to fix this…

sed -i ‘s/nvcaffe_parser/nvparsers/g’ CMakeLists.txt 

Ok…Built! And…

[gstreamer] gstCamera -- attempting to create device v4l2:///dev/video0
[gstreamer] gstCamera -- didn't discover any v4l2 devices
[gstreamer] gstCamera -- device discovery and auto-negotiation failed
[gstreamer] gstCamera -- failed to create device v4l2:///dev/video0
Traceback (most recent call last):
File "my-detection.py", line 31, in
camera = jetson.utils.videoSource("/dev/video0")
Exception: jetson.utils -- failed to create videoSource device

Ok so docker/run.sh also has some other stuff going on, looking for V4L2 devices and such. Ok added device, and it’s working in my nice docker!

sudo docker run -it -p 8888:8888 -p 6006:6006 -p 8265:8265 --rm --runtime nvidia --network host --device /dev/video0 -v /home/chicken/:/dmc nx_setup

So what now… We open ‘/dev/video0’ for V4L2, then replace my inference code below.

img = camera.Capture()     
detections = net.Detect(img)     
display.Render(img)


Ok ... zoom ahead and it's all worked out...
---
After going on the ringer another time: zooming ahead, fixing that jetson issue, involves 
cd jetson-inference/build
make install
ldconfig


---

Saved the sender program at the github 

To get video to the screen

gst-launch-1.0 -v udpsrc port=1234  caps = "application/x-rtp, media=(string)video, clock-rate=(int)90000, encoding-name=(string)H264, payload=(int)96" !  rtph264depay ! decodebin ! videoconvert ! autovideosink

or to a mp4

gst-launch-1.0 -v udpsrc port=1234  caps = "application/x-rtp, media=(string)video, clock-rate=(int)90000, encoding-name=(string)H264, payload=(int)96" !  rtph264depay  ! decodebin ! x264enc ! qtmux ! filesink location=test.mp4 -e
Categories
control envs Gripper Research Locomotion simulation

Meta-world

It is a training task set for 6 DOF (degree-of-freedom) robot arms.

https://github.com/rlworkgroup/metaworld

Here’s a 38MB gif explaining it: https://meta-world.github.io/figures/ml45-1080p.gif

Meta-World is an open-source simulated benchmark for meta-reinforcement learning and multi-task learning consisting of 50 distinct robotic manipulation tasks. We aim to provide task distributions that are sufficiently broad to evaluate meta-RL algorithms’ generalization ability to new behaviors.”

Categories
Behaviour envs meta simulation

Animal-AI 2.0

Like metaworld, but 900 tasks, and with Unity mappings http://animalaiolympics.com/AAI/

github: https://github.com/beyretb/AnimalAI-Olympics

“The Animal-AI Olympics was built using Unity’s ML-Agents Toolkit.

The Python library located in animalai extends ml-agents v0.15.0. Mainly, we add the possibility to change the configuration of arenas between episodes.”

To get an idea of the experiments: http://animalaiolympics.com/AAI/testbed

They had a competition of ‘animal AIs’ in 2019, using EvalAI:

EvalAI

The competition was kindly hosted on EvalAI, an open source web application for AI competitions. Special thanks to Rishabh Jain for his help in setting this up. We will aim to reopen submissions with new hidden files in order to keep some form of competition going.

Deshraj Yadav, Rishabh Jain, Harsh Agrawal, Prithvijit Chattopadhyay, Taranjeet Singh, Akash Jain, Shiv Baran Singh, Stefan Lee and Dhruv Batra (2019) EvalAI: Towards Better Evaluation Systems for AI Agents

arxiv: https://arxiv.org/pdf/1902.03570.pdf

Categories
AI/ML envs Vision

COCO, ShapeNet, Pix3d

These are some examples of data sets, for different reasons

https://cocodataset.org/https://cocodataset.org/

https://www.shapenet.org/about

http://pix3d.csail.mit.edu/

Torchvision:

torchvision.datasets

Categories
3D Research envs simulation

Habitat-Sim

https://github.com/facebookresearch/habitat-sim

A flexible, high-performance 3D simulator with configurable agents, multiple sensors, and generic 3D dataset handling (with built-in support for MatterPort3DGibsonReplica, and other datasets

Categories
AI/ML control deep envs

Flow

https://flow-project.github.io/index.html

Seems like a framework with some sort of traffic environments

Categories
dev envs

Google Football Env

https://github.com/google-research/football
Google made an openAI gym environment for playing football.

It looks better than FIFA.

Categories
AI/ML dev envs neuro simulation

OpenAI Gym MultiNEAT

ok also, just saw this: https://gym.openai.com/evaluations/eval_a0YXWDc4SKeJjyTH7IrHBg/

it doesn’t work apparrently, but could be salvaged into something,

possibly written by this guy https://blog.otoro.net/

https://attentionagent.github.io/ there is no conscious perception of the visual world without attention to it

# Using ES-HyperNEAT to try to solve the Bipedal walker.
# This attempt was not successful. Adjustment of hyperparameters is likely needed.

# A neural network is trained using NeuroEvolution of Augmenting Topologies
# The idea is from the paper: "Evolving Neural Networks through Augmenting Topologies"
# This gist is using MultiNEAT (http://multineat.com/)

import logging
import numpy as np
import pickle

import gym

import MultiNEAT as NEAT

# NEAT setup
params = NEAT.Parameters()
params.PopulationSize = 200;

params.DynamicCompatibility = True;
params.CompatTreshold = 2.0;
params.YoungAgeTreshold = 15;
params.SpeciesMaxStagnation = 100;
params.OldAgeTreshold = 35;
params.MinSpecies = 5;
params.MaxSpecies = 10;
params.RouletteWheelSelection = False;

params.MutateRemLinkProb = 0.02;
params.RecurrentProb = 0;
params.OverallMutationRate = 0.15;
params.MutateAddLinkProb = 0.08;
params.MutateAddNeuronProb = 0.01;
params.MutateWeightsProb = 0.90;
params.MaxWeight = 8.0;
params.WeightMutationMaxPower = 0.2;
params.WeightReplacementMaxPower = 1.0;

params.MutateActivationAProb = 0.0;
params.ActivationAMutationMaxPower = 0.5;
params.MinActivationA = 0.05;
params.MaxActivationA = 6.0;

params.MutateNeuronActivationTypeProb = 0.03;

params.ActivationFunction_SignedSigmoid_Prob = 0.0;
params.ActivationFunction_UnsignedSigmoid_Prob = 0.0;
params.ActivationFunction_Tanh_Prob = 1.0;
params.ActivationFunction_TanhCubic_Prob = 0.0;
params.ActivationFunction_SignedStep_Prob = 1.0;
params.ActivationFunction_UnsignedStep_Prob = 0.0;
params.ActivationFunction_SignedGauss_Prob = 1.0;
params.ActivationFunction_UnsignedGauss_Prob = 0.0;
params.ActivationFunction_Abs_Prob = 0.0;
params.ActivationFunction_SignedSine_Prob = 1.0;
params.ActivationFunction_UnsignedSine_Prob = 0.0;
params.ActivationFunction_Linear_Prob = 1.0;

params.DivisionThreshold = 0.5;
params.VarianceThreshold = 0.03;
params.BandThreshold = 0.3;
params.InitialDepth = 2;
params.MaxDepth = 3;
params.IterationLevel = 1;
params.Leo = False;
params.GeometrySeed = False;
params.LeoSeed = False;
params.LeoThreshold = 0.3;
params.CPPN_Bias = -1.0;
params.Qtree_X = 0.0;
params.Qtree_Y = 0.0;
params.Width = 1.;
params.Height = 1.;
params.Elitism = 0.1;

rng = NEAT.RNG()
rng.TimeSeed()

list = []

for i in range(0,14):
	list.append((-1. +(2.*i/13.), -1., 0.))

for i in range(0,10):
	list.append((-1. +(2.*i/9), -0.5, 0))


substrate = NEAT.Substrate(list,
                           [],
                           [(-1., 1., 0.), (-0.5, 1., 0.), (0.5, 1., 0.), (1., 1., 0.)])

substrate.m_allow_input_hidden_links = False;
substrate.m_allow_input_output_links = False;
substrate.m_allow_hidden_hidden_links = False;
substrate.m_allow_hidden_output_links = False;
substrate.m_allow_output_hidden_links = False;
substrate.m_allow_output_output_links = False;
substrate.m_allow_looped_hidden_links = True;
substrate.m_allow_looped_output_links = False;

substrate.m_allow_input_hidden_links = True;
substrate.m_allow_input_output_links = False;
substrate.m_allow_hidden_output_links = True;
substrate.m_allow_hidden_hidden_links = True;

substrate.m_hidden_nodes_activation = NEAT.ActivationFunction.SIGNED_SIGMOID;
substrate.m_output_nodes_activation = NEAT.ActivationFunction.UNSIGNED_SIGMOID;

substrate.m_with_distance = False;

substrate.m_max_weight_and_bias = 8.0;


def trainNetwork(env, seed):
    # Training parameters
    generationSize = 50
    episode_count = 10
    max_steps = 1000
    # Max reward for environments that reward 1 for each succesfull step (e.g. CartPole-v0)
    max_reward = episode_count * max_steps

    def evaluate(genome):
        net = NEAT.NeuralNetwork()
        genome.BuildESHyperNEATPhenotype(net, substrate, params)

        cum_reward = 0

        for i in xrange(episode_count):
            ob = env.reset()
            net.Flush()

            for j in xrange(max_steps):
                # get next action
                net.Input(ob)
                net.Activate()
                o = net.Output()
                action = np.clip(o,-1,1)
                ob, reward, done, _ = env.step(action)
                cum_reward += reward
                if done:
                    break

        return cum_reward

    # Create initial genome
    g = NEAT.Genome(0, 24, 0, 4, False, 
                    NEAT.ActivationFunction.TANH, NEAT.ActivationFunction.TANH, 0, params)
    pop = NEAT.Population(g, params, True, 1.0, seed)

    current_best = None

    for generation in range(generationSize):
        for i_episode, genome in enumerate(NEAT.GetGenomeList(pop)):
            reward = evaluate(genome)

            if reward == max_reward:
                return pickle.dumps(genome)

            genome.SetFitness(reward)

        print('Generation: {}, max fitness: {}'.format(generation,
                            max((x.GetFitness() for x in NEAT.GetGenomeList(pop)))))
        current_best = pickle.dumps(pop.GetBestGenome())
        pop.Epoch()


    return current_best

env_name = "BipedalWalker"

if __name__ == '__main__':
    # Test the algorithm multiple times
    for test_case in xrange(0, 1):
        # setup logger, environment and monitor
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        env = gym.make("%s-v2" % env_name)
        outdir = "/tmp/neat-%s-results-%d" % (env_name, test_case)
        env.monitor.start(outdir, force=True)

        # Train network
        learned = trainNetwork(env, test_case)

        # Test trained network on 1000 episodes
        learned_genome = pickle.loads(learned)
        net = NEAT.NeuralNetwork()
        learned_genome.BuildESHyperNEATPhenotype( net,substrate, params)

        episode_count = 1000
        max_steps = 1000

        for i in xrange(episode_count):
            ob = env.reset()
            net.Flush()

            for j in xrange(max_steps):
                # get next action
                net.Input(ob)
                net.Activate()
                o = net.Output()
                action = np.clip(o,-1,1)
                ob, reward, done, _ = env.step(action)
                if done:
                    break


        # Dump result info to disk
        env.monitor.close()