Upload 4 files

config.json

@@ -0,0 +1,19 @@
+{
+  "architectures": [
+    "MrPongForRL"
+  ],
+  "model_type": "mr_pong",
+  "obs_dim": 12,
+  "action_dim": 3,
+  "hidden_dims": [
+    160,
+    160
+  ],
+  "activation": "tanh",
+  "torch_dtype": "float32",
+  "global_step": 9975808,
+  "auto_map": {
+    "AutoConfig": "configuration_mrs_paleta.MrPongConfig",
+    "AutoModel": "modeling_mr_pong.MrPongForRL"
+  }
+}

configuration_mr_pong.py

@@ -0,0 +1,39 @@
+"""
+Configuration class for Mrs. Paleta Ping Pong RL Agent.
+"""
+
+from transformers import PretrainedConfig
+
+
+class MrPongConfig(PretrainedConfig):
+    model_type = "mrs_paleta"
+
+    def __init__(
+        self,
+        obs_dim: int = 12,
+        action_dim: int = 3,
+        hidden_dims: list = None,
+        activation: str = "tanh",
+        table_width: float = 800.0,
+        table_height: float = 500.0,
+        paddle_width: float = 14.0,
+        paddle_height: float = 80.0,
+        paddle_speed: float = 8.0,
+        ball_speed_max: float = 16.0,
+        max_rally_steps: int = 1500,
+        **kwargs
+    ):
+        if hidden_dims is None:
+            hidden_dims = [160, 160]
+        self.obs_dim = obs_dim
+        self.action_dim = action_dim
+        self.hidden_dims = hidden_dims
+        self.activation = activation
+        self.table_width = table_width
+        self.table_height = table_height
+        self.paddle_width = paddle_width
+        self.paddle_height = paddle_height
+        self.paddle_speed = paddle_speed
+        self.ball_speed_max = ball_speed_max
+        self.max_rally_steps = max_rally_steps
+        super().__init__(**kwargs)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dae8d682e24a85661abd97f1d6ba09602aea5e395c91a78f9cdb1844ca8d161
+size 114536

modeling_mr_pong.py

@@ -0,0 +1,109 @@
+"""
+PyTorch Modeling class for Mrs. Paleta Ping Pong RL Agent.
+Compatible with Hugging Face AutoModel via trust_remote_code=True.
+"""
+
+from typing import Optional, Tuple, Union, Dict, Any
+from dataclasses import dataclass
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+try:
+    from .configuration_mr_pong import MrPongConfig
+except ImportError:
+    from configuration_mr_pong import MrPongConfig
+
+
+@dataclass
+class MrPongOutput(ModelOutput):
+    """
+    Model output for Mr Pong.
+    """
+    logits: torch.FloatTensor = None
+    value_estimate: Optional[torch.FloatTensor] = None
+    action: Optional[torch.LongTensor] = None
+
+
+class MrPongForRL(PreTrainedModel):
+    config_class = MrPongConfig
+    base_model_prefix = "mr_pong"
+
+    def __init__(self, config: MrPongConfig):
+        super().__init__(config)
+        self.config = config
+
+        act_fn = nn.Tanh if config.activation == "tanh" else (nn.ReLU if config.activation == "relu" else nn.GELU)
+
+        layers = []
+        in_dim = config.obs_dim
+        for h_dim in config.hidden_dims:
+            layers.append(nn.Linear(in_dim, h_dim))
+            layers.append(act_fn())
+            in_dim = h_dim
+
+        self.trunk = nn.Sequential(*layers)
+        self.actor = nn.Linear(in_dim, config.action_dim)
+        self.critic = nn.Linear(in_dim, 1)
+
+        self.post_init()
+
+    def forward(
+        self,
+        observation: torch.FloatTensor,
+        deterministic: bool = True,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple[torch.FloatTensor, torch.FloatTensor], MrPongOutput]:
+        """
+        Forward pass returning action logits, state-value estimates, and greedily/sampled chosen action.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not isinstance(observation, torch.Tensor):
+            observation = torch.tensor(observation, dtype=torch.float32)
+
+        if observation.ndim == 1:
+            observation = observation.unsqueeze(0)
+
+        # Slice or pad to expected input dimension
+        if observation.shape[-1] > self.config.obs_dim:
+            observation = observation[..., :self.config.obs_dim]
+        elif observation.shape[-1] < self.config.obs_dim:
+            pad_size = self.config.obs_dim - observation.shape[-1]
+            observation = nn.functional.pad(observation, (0, pad_size))
+
+        features = self.trunk(observation)
+        logits = self.actor(features)
+        value = self.critic(features)
+
+        if deterministic:
+            action = torch.argmax(logits, dim=-1)
+        else:
+            dist = torch.distributions.Categorical(logits=logits)
+            action = dist.sample()
+
+        if not return_dict:
+            return logits, value, action
+
+        return MrPongOutput(
+            logits=logits,
+            value_estimate=value,
+            action=action
+        )
+
+    @torch.no_grad()
+    def act(self, observation: Union[np.ndarray, list, torch.Tensor], deterministic: bool = True) -> int:
+        """
+        High-level inference method returning single integer action (0: Stay, 1: Up, 2: Down).
+        """
+        self.eval()
+        if not isinstance(observation, torch.Tensor):
+            observation = torch.tensor(observation, dtype=torch.float32, device=self.device)
+        else:
+            observation = observation.to(self.device)
+
+        out = self.forward(observation, deterministic=deterministic, return_dict=True)
+        return out.action.squeeze().item()