Fix ActorCritic family (#133)

* fix reinforce val loss calc as mean

* make ac loss coef calc consistent

* update val_loss_coef to 0.01 by default

* add tmp ac config to run

* remove tmp config

* add default coef values to prevent inheritance breakage

* add val loss coef to PPO

* update ac family search range

* add entropy

* reinforce back to sum per textbook

* monitor log 4 digits loss

* Changing sum to mean in reinforce policy loss calc

* Update fn to new pyTorch syntax

* Fixing computation graph error

* remove unnecessary mean from reinforce loss

* add assert_trained method

* fix actorcritic class

* expand sil param, check compute carry gradient, tuning needed

* standardize args, restore policy_loss_coef to ac calc

* compact rewrite of reinforce loss works

* compact rewrite of actor critic loss works

* fix PPO

* use same rewrite for sil log_probs, still need fixing/tuning

* add screen to linux plot

* fix platform

* split cmd

* use xvfb wrapper

* debug log

* save as context manager

* add xvfb, prevent linux plot crash now

* fix silly SIL loss var typo. working, need turning

* add sha to all spec written

* update orca install, use npm. update npm build

* fix plotting for linux

Author: kengz

.circleci/config.yml

@@ -24,7 +24,7 @@ jobs:
           name: Install apt-get packages for lab
           command: |
             sudo apt-get update
-            sudo apt-get install -y python-numpy python-dev cmake zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6
+            sudo apt-get install -y python-numpy python-dev cmake libhdf5-dev libopenblas-dev zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6
           environment:
             LD_PRELOAD: /usr/lib/libtcmalloc_minimal.so.4
 
@@ -34,10 +34,9 @@ jobs:
             if which yarn >/dev/null; then
               echo "Yarn is already installed"
             else
-              sudo npm install -g yarn
+              sudo npm install --unsafe-perm=true --allow-root -g yarn [email protected] orca
             fi
             yarn install
-            yarn global add [email protected] orca
       - save_cache:
           paths:
             - node_modules
@@ -100,7 +99,7 @@ jobs:
           name: Install apt-get packages for lab
           command: |
             sudo apt-get update
-            sudo apt-get install -y python-numpy python-dev cmake zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6
+            sudo apt-get install -y python-numpy python-dev cmake libhdf5-dev libopenblas-dev zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6
       - run: echo '. ~/miniconda3/etc/profile.d/conda.sh' >> $BASH_ENV
       - run:
           name: Run Python tests

bin/setup_macOS

@@ -23,32 +23,20 @@ for item in "${hb_list[@]}"; do
   brew info "${item}" | grep --quiet "Not installed" && brew install "${item}"
 done
 
-echo "--- Installing Atom and Hydrogen for interactive computing ---"
-if which atom >/dev/null; then
-  echo "Atom is already installed"
-else
-  brew cask install atom
-fi
-if apm ls | grep " Hydrogen" >/dev/null; then
-  echo "Hydrogen is already installed"
-else
-  apm install hydrogen
-fi
-
 echo "--- Installing NodeJS Lab interface ---"
 if which node >/dev/null; then
   echo "NodeJS is already installed"
 else
   brew install node
   brew install yarn
+  npm install --unsafe-perm=true --allow-root -g [email protected] orca
 fi
 
 echo "--- Installing npm modules for Lab interface ---"
 if [ -d ./node_modules ]; then
   echo "Npm modules are already installed"
 else
   yarn install
-  yarn global add [email protected] orca
 fi
 
 echo "--- Installing Python for Lab backend ---"

bin/setup_ubuntu

@@ -7,39 +7,24 @@ trap "exit" INT
 
 echo "--- Installing system dependencies ---"
 sudo add-apt-repository -y ppa:ubuntu-toolchain-r/test
-sudo add-apt-repository -y ppa:webupd8team/atom
 sudo apt-get update
-sudo apt-get install -y cmake gcc-4.9 g++-4.9 git
-sudo apt-get install -y libhdf5-dev libopenblas-dev
-sudo apt-get install -y cmake zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6
-
-echo "--- Installing Atom and Hydrogen for interactive computing ---"
-if which atom >/dev/null; then
-  echo "Atom is already installed"
-else
-  sudo apt-get install -y atom
-fi
-if apm ls | grep " Hydrogen" >/dev/null; then
-  echo "Hydrogen is already installed"
-else
-  apm install hydrogen
-fi
+sudo apt-get install -y git cmake gcc g++
+sudo apt-get install -y zlib1g-dev libjpeg-dev xvfb libav-tools xorg-dev python-opengl libboost-all-dev libsdl2-dev swig build-essential libstdc++6 libhdf5-dev libopenblas-dev
 
 echo "--- Installing NodeJS Lab interface ---"
 if which node >/dev/null; then
   echo "Nodejs is already installed"
 else
   curl -sL https://deb.nodesource.com/setup_8.x | sudo -E bash -
   sudo apt-get install -y nodejs
-  sudo npm install -g yarn
+  sudo npm install --unsafe-perm=true --allow-root -g yarn [email protected] orca
 fi
 
 echo "--- Installing npm modules for Lab interface ---"
 if [ -d ./node_modules ]; then
   echo "Npm modules are already installed"
 else
   yarn install
-  yarn global add [email protected] orca
 fi
 
 echo "--- Installing Python for Lab backend ---"

slm_lab/agent/algorithm/actor_critic.py

@@ -61,7 +61,7 @@ class ActorCritic(Reinforce):
         "add_entropy": false,
         "entropy_coef": 0.01,
         "policy_loss_coef": 1.0,
-        "val_loss_coef": 1.0,
+        "val_loss_coef": 0.01,
         "continuous_action_clip": 2.0,
         "training_frequency": 1,
         "training_epoch": 8
@@ -87,6 +87,8 @@ def init_algorithm_params(self):
             explore_var_start=np.nan,
             explore_var_end=np.nan,
             explore_anneal_epi=np.nan,
+            policy_loss_coef=1.0,
+            val_loss_coef=1.0,
         ))
         util.set_attr(self, self.algorithm_spec, [
             'action_pdtype',
@@ -263,9 +265,9 @@ def train_shared(self):
             batch = self.sample()
             with torch.no_grad():
                 advs, v_targets = self.calc_advs_v_targets(batch)
-            policy_loss = self.calc_policy_loss(advs)  # from actor
+            policy_loss = self.calc_policy_loss(batch, advs)  # from actor
             val_loss = self.calc_val_loss(batch, v_targets)  # from critic
-            loss = self.policy_loss_coef * policy_loss + self.val_loss_coef * val_loss
+            loss = policy_loss + val_loss
             self.net.training_step(loss=loss)
             # reset
             self.to_train = 0
@@ -282,9 +284,7 @@ def train_separate(self):
         '''
         if self.to_train == 1:
             batch = self.sample()
-            with torch.no_grad():
-                advs, v_targets = self.calc_advs_v_targets(batch)
-            policy_loss = self.train_actor(advs)
+            policy_loss = self.train_actor(batch)
             val_loss = self.train_critic(batch)
             loss = val_loss + abs(policy_loss)
             # reset
@@ -295,9 +295,11 @@ def train_separate(self):
             self.last_loss = loss.item()
         return self.last_loss
 
-    def train_actor(self, advs):
+    def train_actor(self, batch):
         '''Trains the actor when the actor and critic are separate networks'''
-        policy_loss = self.calc_policy_loss(advs)
+        with torch.no_grad():
+            advs, _v_targets = self.calc_advs_v_targets(batch)
+        policy_loss = self.calc_policy_loss(batch, advs)
         self.net.training_step(loss=policy_loss)
         return policy_loss
 
@@ -314,15 +316,14 @@ def train_critic(self, batch):
         val_loss = total_val_loss / self.training_epoch
         return val_loss
 
-    def calc_policy_loss(self, advs):
+    def calc_policy_loss(self, batch, advs):
         '''Calculate the actor's policy loss'''
         assert len(self.body.log_probs) == len(advs), f'{len(self.body.log_probs)} vs {len(advs)}'
-        log_probs = torch.tensor(self.body.log_probs, requires_grad=True)
-        entropies = torch.tensor(self.body.entropies, requires_grad=True)
+        log_probs = torch.stack(self.body.log_probs)
+        policy_loss = - self.policy_loss_coef * log_probs * advs
         if self.add_entropy:
-            policy_loss = (- log_probs * advs) - self.entropy_coef * entropies
-        else:
-            policy_loss = - log_probs * advs
+            entropies = torch.stack(self.body.entropies)
+            policy_loss += (-self.entropy_coef * entropies)
         policy_loss = torch.mean(policy_loss)
         if torch.cuda.is_available() and self.net.gpu:
             policy_loss = policy_loss.cuda()
@@ -334,7 +335,7 @@ def calc_val_loss(self, batch, v_targets):
         v_targets = v_targets.unsqueeze(dim=-1)
         v_preds = self.calc_v(batch['states'], evaluate=False).unsqueeze_(dim=-1)
         assert v_preds.shape == v_targets.shape
-        val_loss = self.net.loss_fn(v_preds, v_targets)
+        val_loss = self.val_loss_coef * self.net.loss_fn(v_preds, v_targets)
         if torch.cuda.is_available() and self.net.gpu:
             val_loss = val_loss.cuda()
         logger.debug(f'Critic value loss: {val_loss:.2f}')

slm_lab/agent/algorithm/ppo.py

@@ -67,6 +67,7 @@ def init_algorithm_params(self):
             explore_var_start=np.nan,
             explore_var_end=np.nan,
             explore_anneal_epi=np.nan,
+            val_loss_coef=1.0,
         ))
         util.set_attr(self, self.algorithm_spec, [
             'action_pdtype',
@@ -80,6 +81,7 @@ def init_algorithm_params(self):
             'lam',
             'clip_eps',
             'entropy_coef',
+            'val_loss_coef',
             'training_frequency',  # horizon
             'training_epoch',
         ])
@@ -90,6 +92,8 @@ def init_algorithm_params(self):
         self.action_policy_update = getattr(policy_util, self.action_policy_update)
         for body in self.agent.nanflat_body_a:
             body.explore_var = self.explore_var_start
+        # PPO uses GAE
+        self.calc_advs_v_targets = self.calc_gae_advs_v_targets
 
     @lab_api
     def init_nets(self):
@@ -111,20 +115,20 @@ def calc_log_probs(self, batch, use_old_net=False):
         # get ActionPD, don't append to state_buffer
         ActionPD, _pdparam, _body = policy_util.init_action_pd(states[0].cpu().numpy(), self, self.body, append=False)
         # construct log_probs for each state-action
-        pdparams = self.calc_pdparam(states)
+        pdparams = self.calc_pdparam(states, evaluate=False)
         log_probs = []
         for idx, pdparam in enumerate(pdparams):
             _action, action_pd = policy_util.sample_action_pd(ActionPD, pdparam, self.body)
             log_prob = action_pd.log_prob(actions[idx])
             log_probs.append(log_prob)
-        log_probs = torch.tensor(log_probs)
+        log_probs = torch.stack(log_probs)
         if use_old_net:
             # swap back
             self.old_net = self.net
             self.net = self.tmp_net
         return log_probs
 
-    def calc_loss(self, batch):
+    def calc_policy_loss(self, batch, advs):
         '''
         The PPO loss function (subscript t is omitted)
         L^{CLIP+VF+S} = E[ L^CLIP - c1 * L^VF + c2 * S[pi](s) ]
@@ -133,35 +137,34 @@ def calc_loss(self, batch):
         1. L^CLIP = E[ min(ratio * A, clip(ratio, 1-eps, 1+eps) * A) ]
         where ratio = pi(a|s) / pi_old(a|s)
 
-        2. L^VF = E[ (V(s_t) - V^target)^2 ]
+        2. L^VF = E[ mse(V(s_t), V^target) ]
 
         3. S = E[ entropy ]
         '''
         # decay clip_eps by episode
         clip_eps = policy_util._linear_decay(self.clip_eps, 0.1 * self.clip_eps, self.clip_eps_anneal_epi, self.body.env.clock.get('epi'))
 
-        with torch.no_grad():
-            adv_targets, v_targets = self.calc_gae_advs_v_targets(batch)
-
         # L^CLIP
         log_probs = self.calc_log_probs(batch, use_old_net=False)
         old_log_probs = self.calc_log_probs(batch, use_old_net=True)
         assert log_probs.shape == old_log_probs.shape
-        assert adv_targets.shape == log_probs.shape
+        assert advs.shape == log_probs.shape
         ratios = torch.exp(log_probs - old_log_probs)
-        sur_1 = ratios * adv_targets
-        sur_2 = torch.clamp(ratios, 1.0 - clip_eps, 1.0 + clip_eps) * adv_targets
+        sur_1 = ratios * advs
+        sur_2 = torch.clamp(ratios, 1.0 - clip_eps, 1.0 + clip_eps) * advs
         # flip sign because need to maximize
         clip_loss = -torch.mean(torch.min(sur_1, sur_2))
 
-        # L^VF
-        val_loss = self.calc_val_loss(batch, v_targets)  # from critic
+        # L^VF (inherit from ActorCritic)
 
         # S entropy bonus
-        ent_mean = torch.mean(torch.tensor(self.body.entropies))
-        ent_penalty = -self.entropy_coef * ent_mean
-        loss = clip_loss + val_loss + ent_penalty
-        return loss
+        ent_penalty = 0
+        for e in self.body.entropies:
+            ent_penalty += (-self.entropy_coef * e)
+        ent_penalty /= len(self.body.entropies)
+
+        policy_loss = clip_loss + ent_penalty
+        return policy_loss
 
     def train_shared(self):
         '''
@@ -171,8 +174,13 @@ def train_shared(self):
             batch = self.sample()
             total_loss = torch.tensor(0.0)
             for _ in range(self.training_epoch):
-                loss = self.calc_loss(batch)
-                self.net.training_step(loss=loss)
+                with torch.no_grad():
+                    advs, v_targets = self.calc_advs_v_targets(batch)
+                policy_loss = self.calc_policy_loss(batch, advs)  # from actor
+                val_loss = self.calc_val_loss(batch, v_targets)  # from critic
+                loss = policy_loss + val_loss
+                # retain for entropies etc.
+                self.net.training_step(loss=loss, retain_graph=True)
                 total_loss += loss.cpu()
             loss = total_loss / self.training_epoch
             net_util.copy(self.net, self.old_net)
@@ -190,15 +198,9 @@ def train_separate(self):
         '''
         if self.to_train == 1:
             batch = self.sample()
-            total_loss = torch.tensor(0.0)
-            for _ in range(self.training_epoch):
-                loss = self.calc_loss(batch)
-                # to reuse loss for critic
-                self.net.training_step(loss=loss, retain_graph=True)
-                # critic.optim.step using the same loss
-                self.critic.training_step(loss=loss)
-                total_loss += loss.cpu()
-            loss = total_loss / self.training_epoch
+            policy_loss = self.train_actor(batch)
+            val_loss = self.train_critic(batch)
+            loss = val_loss + abs(policy_loss)
             net_util.copy(self.net, self.old_net)
             net_util.copy(self.critic, self.old_critic)
             # reset
@@ -208,3 +210,15 @@ def train_separate(self):
             logger.debug(f'Loss: {loss:.2f}')
             self.last_loss = loss.item()
         return self.last_loss
+
+    def train_actor(self, batch):
+        '''Trains the actor when the actor and critic are separate networks'''
+        total_policy_loss = torch.tensor(0.0)
+        for _ in range(self.training_epoch):
+            with torch.no_grad():
+                advs, _v_targets = self.calc_advs_v_targets(batch)
+            policy_loss = self.calc_policy_loss(batch, advs)
+            # retain for entropies etc.
+            self.net.training_step(loss=policy_loss, retain_graph=True)
+        val_loss = total_policy_loss / self.training_epoch
+        return policy_loss

slm_lab/agent/algorithm/reinforce.py

@@ -159,15 +159,14 @@ def calc_policy_loss(self, batch):
         adv_std += 1e-08
         advs = (advs - advs.mean()) / adv_std
         assert len(self.body.log_probs) == len(advs), f'{len(self.body.log_probs)} vs {len(advs)}'
-        policy_loss = torch.tensor(0.0)
+        log_probs = torch.stack(self.body.log_probs)
+        policy_loss = - log_probs * advs
+        if self.add_entropy:
+            entropies = torch.stack(self.body.entropies)
+            policy_loss += (-self.entropy_coef * entropies)
+        policy_loss = torch.sum(policy_loss)
         if torch.cuda.is_available() and self.net.gpu:
-            advs = advs.cuda()
             policy_loss = policy_loss.cuda()
-        for logp, adv, ent in zip(self.body.log_probs, advs, self.body.entropies):
-            if self.add_entropy:
-                policy_loss += (-logp * adv - self.entropy_coef * ent).cpu()
-            else:
-                policy_loss += (-logp * adv).cpu()
         return policy_loss
 
     @lab_api