Demo: SLM Lab as pip module for lightweight usecases

slm-lab-pip-install-usage.py

'''
Demo: SLM Lab as pip module for lightweight usecases

Installation:

1. Clone SLM-Lab
```
git clone https://github.com/kengz/SLM-Lab.git
cd SLM-Lab
```

2. Install node, yarn, electron and orca for the environment and analysis module
```
brew install node yarn
yarn install
```

And if you don't already have it, install Conda for clean dependency management
```
curl -O https://repo.continuum.io/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
bash Miniconda3-latest-MacOSX-x86_64.sh -b
rm Miniconda3-latest-MacOSX-x86_64.sh
echo '. ~/miniconda3/etc/profile.d/conda.sh' >> ~/.bash_profile
source ~/.bash_profile
```

3. Install SLM Lab as pip module
```
conda create -n lab python=3.6 ipykernel -y
conda activate lab
conda env update -f environment.yml
pip install -e .
```

Now it is ready for usage.
'''
import os
# NOTE increase if needed. Pytorch thread overusage https://github.com/pytorch/pytorch/issues/975
os.environ['OMP_NUM_THREADS'] = '1'
from slm_lab.agent import Agent
from slm_lab.env import OpenAIEnv
from slm_lab.experiment import analysis
from slm_lab.experiment.monitor import Body, InfoSpace, enable_aeb_space
from slm_lab.lib import logger, util
from slm_lab.spec import spec_util


class Session:
    '''The class which initializes the agent, environment, and runs them.'''

    def __init__(self, spec, info_space):
        self.spec = spec
        self.info_space = info_space
        self.index = self.info_space.get('session')

        # init singleton agent and env
        self.env = OpenAIEnv(self.spec)
        body = Body(self.env, self.spec['agent'])
        self.agent = Agent(self.spec, self.info_space, body=body)

        enable_aeb_space(self)  # to use lab's data analysis framework
        logger.info(f'Initialized session {self.index}')

    def run_episode(self):
        self.env.clock.tick('epi')
        reward, state, done = self.env.reset()
        self.agent.reset(state)
        while not done:
            self.env.clock.tick('t')
            action = self.agent.act(state)
            reward, state, done = self.env.step(action)
            self.agent.update(action, reward, state, done)
        self.agent.body.log_summary()

    def close(self):
        self.agent.close()
        self.env.close()
        logger.info('Session done and closed.')

    def run(self):
        while self.env.clock.get('epi') <= self.env.max_episode:
            self.run_episode()
        self.data = analysis.analyze_session(self)  # session fitness
        self.close()
        return self.data


# To use SLM-Lab's existing spec. Alternatively, you can write one yourself too
spec = spec_util.get(spec_file='ppo.json', spec_name='ppo_mlp_shared_cartpole')
info_space = InfoSpace()
# set proper env variables for the lab
os.environ['PREPATH'] = util.get_prepath(spec, info_space)
os.environ['lab_mode'] = 'dev'  # set to 'train' to run at full speed

# inspect the agent spec; edit if you wish to
print(spec['agent'])
# edit the env spec to run for less episodes
spec['env'][0]['max_episode'] = 100

# initialize and run session
sess = Session(spec, info_space)
data = sess.run()

print(f'Data is available at {util.smart_path(os.environ["PREPATH"])}')