RILIR 复现 & 一些 idea

伪代码：

在 if done 的时候，在环境中已经跑了一个 trajectory 了，利用当前的 trajectory 和专家的 demo 求一下 reward（文章中用的是 optimal transport 的几种方法）
否则，就继续在 observation 的基础上利用 actor 学到的策略 sample 出 action，并用 list 记录下当前的 \((o_t, a_t, r_t)\) 作为 trajectory 的一部分。

整体来看：
首先整体的强化学习框架是 actor-critic，其中 actor 输出 action（的分布），更新使用 critic，而 critic 作为 actor 的评估函数，更新使用时序残差的 MSE，即
这一项。actor 如何更新的 （待填）。

IL 中的重要问题，如何根据专家数据和 actor 生成之间的差距来获取 reward？在 RILIR 中，采用的是 trajectory 级别的模仿学习，即将当前的 trajectory 和 expert demo trajectory 拿出来做一个 OT（理解成相似程度），得到 reward

actor 实现的时候有个小细节：先用 trunk 函数得到一个 embedding，将 input_dim 嵌入到 feature_dim，这个函数之后再将 feather_dim 嵌入到 hidden_dim 和最后的结果维度。

这个 loss 函数是用来更新 critic 的。（actor-critic 本身就是 决策网络-价值网络），而更新 actor 就按照正常的方式：

inv_dyn 为 \(f_\theta()\)，为 inverse dynamic，用来计算预测出的两个 observation 之间的预测 action。
critic 或者 critic_target 即为 \(Q\) function，用于求 \(Q(s,a)\)
encoder 即为 \(\phi\)，RILIR 中用 CNN 实现的。

	def update_critic(self, obs, action, reward, discount, next_obs, step, next_obs_grad,
                   oe_t, oe_t1, ae_t):
		metrics = dict()

		with torch.no_grad():
			stddev = utils.schedule(self.stddev_schedule, step)
			dist = self.actor(next_obs, stddev)
			next_action = dist.sample(clip=self.stddev_clip)
			target_Q1, target_Q2 = self.critic_target(next_obs, next_action)
			target_V = torch.min(target_Q1, target_Q2)
			target_Q = reward + (discount * target_V)

		Q1, Q2 = self.critic(obs, action)

		critic_loss = F.mse_loss(Q1, target_Q) + F.mse_loss(Q2, target_Q)
  
		# calculate loss for inverse dynamics
		obs_new = torch.cat([obs, oe_t], dim=0)
		obs_new1 = torch.cat([next_obs_grad, oe_t1], dim=0)
		action_new = torch.cat([action, ae_t], dim=0)
  
		obs_embedding = self.critic.trunk(obs_new)
		next_obs_embedding = self.critic.trunk(obs_new1)
		inv_dist = self.inv_dyn(obs_embedding, next_obs_embedding, stddev)
		if self.log_like: # log_like = False
			log_prob_inv = inv_dist.log_prob(action).sum(-1, keepdim=True)
			inv_loss = - log_prob_inv.mean() * self.inv_weight
		else:
			# use l2 loss for inverse dynamics
			action_pred = inv_dist.sample(clip=self.stddev_clip)
			inv_loss = F.mse_loss(action_pred, action_new) * self.inv_weight
  
		critic_loss += inv_loss


		# optimize encoder and critic
		if self.use_encoder:
			self.encoder_opt.zero_grad(set_to_none=True)
			self.dyn_opt.zero_grad(set_to_none=True)
		self.critic_opt.zero_grad(set_to_none=True)
		critic_loss.backward()
		self.critic_opt.step()
		if self.use_encoder:
			self.encoder_opt.step()
			self.dyn_opt.step()

		if self.use_tb:
			metrics['critic_target_q'] = target_Q.mean().item()
			metrics['critic_q1'] = Q1.mean().item()
			metrics['critic_q2'] = Q2.mean().item()
			metrics['critic_loss'] = critic_loss.item()
			metrics["inv_dyn_loss"] = inv_loss.item()
			metrics["inv_weights"] = self.inv_weight
			
		return metrics

在 discriminator 中，训练一个网络来分辨 expert 和自己的，具体的，将 [expert_obs, policy_obs] 和 [expert_action, policy_action] 并一起（原程序中转置了一下）和对应的 [0.., 1..] 求 cross entropy，对 loss 反向传播即可。