Update on "[DRAFT] ppo chess with llm and ConditionalPolicySwitch to sunfish bot"

[ghstack-poisoned]

Author: mikaylagawarecki

examples/agents/ppo-chess-llm.py

@@ -22,7 +22,7 @@
 from torchrl.collectors import SyncDataCollector
 from torchrl.data import NonTensor
 from torchrl.data.replay_buffers.samplers import SliceSamplerWithoutReplacement
-from torchrl.data.tensor_specs import Composite
+from torchrl.data.tensor_specs import Box, Composite, TensorSpec
 from torchrl.envs import ChessEnv
 from torchrl.envs.transforms import (
  ConditionalPolicySwitch,
@@ -75,6 +75,41 @@ def _reset(self, tensordict, tensordict_reset):
  return tensordict_reset
 
 
+class Score(Transform):
+ def __init__(self, input_queue, output_queue):
+ super().__init__()
+ self.input_queue = input_queue
+ self.output_queue = output_queue
+
+ def _step(self, tensordict, next_tensordict):
+ fen = next_tensordict["fen"]
+ self.input_queue.put(fen)
+ _, score = self.output_queue.get()
+ next_tensordict["score"] = torch.tensor(
+ score, device="cuda:7", dtype=torch.bfloat16
+ )
+ return next_tensordict
+
+ def _reset(self, tensordict, tensordict_reset):
+ fen = tensordict_reset["fen"]
+ self.input_queue.put(fen)
+ _, score = self.output_queue.get()
+ tensordict_reset["score"] = torch.tensor(
+ score, device="cuda:7", dtype=torch.bfloat16
+ )
+ return tensordict_reset
+
+ def transform_observation_spec(self, observation_spec: Composite):
+ if not isinstance(observation_spec, Composite):
+ raise ValueError(
+ f"observation_spec was expected to be of type Composite. Got {type(observation_spec)} instead."
+ )
+ observation_spec["observation"] = TensorSpec(
+ (), Box(), dtype=torch.bfloat16, device="cuda:7"
+ )
+ return observation_spec
+
+
 class LLMInputTransform(Transform):
  def __init__(self, san_moves):
  super().__init__()
@@ -159,9 +194,14 @@ def run_player(input_queue, output_queue):
 
  output = process.stdout.readline()
  if output:
- # print(f"Output: {output.strip()}")
+ print(f"Output: {output.strip()}")
  move = re.search(r"bestmove (.*)", output.strip()).group(1)
- output_queue.put(move)
+
+ output = process.stdout.readline()
+ print(f"Output scores: {output.strip()}")
+ score = re.search(r"score (.*)", output.strip()).group(1)
+
+ output_queue.put((move, int(score)))
 
  except queue.Empty:
  continue
@@ -179,7 +219,7 @@ def run_player(input_queue, output_queue):
 def setup_env(input_queue, output_queue, tokenizer):
  def policy_sunfish(td):
  input_queue.put(td["fen"])
- move = output_queue.get()
+ move, _ = output_queue.get()
  san = env.board.san(chess.Move.from_uci(move))
  san_idx = env.san_moves.index(san)
  td["action"] = torch.tensor(san_idx)
@@ -205,6 +245,7 @@ def policy_sunfish(td):
  tokenizer=tokenizer,
  )
  )
+ env.append_transform(Score(input_queue, output_queue))
  env.reset()
  return env
 
@@ -405,16 +446,28 @@ def remove_logits(td):
  return_composite=True,
  )
 
- class CriticHead(torch.nn.Module):
+ # class CriticHead(torch.nn.Module):
+ # def __init__(self):
+ # super().__init__()
+ # self.m = torch.nn.Linear(3584, 1, dtype=torch.bfloat16)
+
+ # def forward(self, hidden):
+ # return self.m(hidden).squeeze(-1).sum(-1, keepdim=True)
+
+ # critic_llm_policy = Seq(
+ # Mod(CriticHead(), in_keys=["hidden"], out_keys=["state_value"]),
+ # )
+
+ class CriticLLMPolicy(torch.nn.Module):
  def __init__(self):
  super().__init__()
- self.m = torch.nn.Linear(3584, 1, dtype=torch.bfloat16)
 
- def forward(self, hidden):
- return self.m(hidden).squeeze(-1).sum(-1, keepdim=True)
+ def forward(self, score):
+ # breakpoint()
+ return score.unsqueeze(-1)
 
  critic_llm_policy = Seq(
- Mod(CriticHead(), in_keys=["hidden"], out_keys=["state_value"]),
+ Mod(CriticLLMPolicy(), in_keys=["score"], out_keys=["state_value"]),
  )
 
  return actor_llm_policy, data_llm_policy, critic_llm_policy, tokenizer
@@ -425,12 +478,21 @@ def play(env, data_llm_policy, actor_llm_policy, tokenizer):
 
  rb = ReplayBuffer(
  storage=LazyStackStorage(100),
- batch_size=8,
+ batch_size=48,
  sampler=SliceSamplerWithoutReplacement(slice_len=8, end_key=("next", "done")),
  )
 
+ # def breakpointy(td):
+ # breakpoint()
+ # return td
+
+ # rb.append_transform(breakpointy)
+
+ # Temporarily patched fbcode/pytorch/tensordict/tensordict/_lazy.py?lines=1502
  rb.append_transform(lambda td: td.densify(layout=torch.jagged))
 
+ # rb.append_transform(breakpointy)
+
  # obs_tokens in layout=torch.jagged errors with Qwen
  # File "/home/mg1998/.conda/envs/rl/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 859, in forward
  # cache_position = torch.arange(
@@ -486,6 +548,7 @@ def obs_token_transform(td):
 
  data = gae(data)
  loss = loss_module(data)
+ breakpoint()
  loss.sum(reduce=True).backward()
  torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 0.5)
  optim.step()