diff --git a/src/lerobot/policies/wall_x/modeling_wall_x.py b/src/lerobot/policies/wall_x/modeling_wall_x.py
index d0c192fb4..a99995bef 100644
--- a/src/lerobot/policies/wall_x/modeling_wall_x.py
+++ b/src/lerobot/policies/wall_x/modeling_wall_x.py
@@ -34,14 +34,10 @@ lerobot-train \
 ```
 """
 
-import builtins
-import glob
+
 import math
-import os
 from os import PathLike
-import sys
 from collections import deque
-from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 from PIL import Image
 
@@ -50,13 +46,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from peft import LoraConfig, get_peft_model
-from safetensors.torch import load_file
 from torch import Tensor
 from torch.distributions import Beta
 from torch.nn import CrossEntropyLoss
 from torchdiffeq import odeint
-from transformers import AutoConfig, AutoProcessor
-from transformers.activations import ACT2FN
+from transformers import AutoProcessor
 from transformers.cache_utils import (
     StaticCache,
 )
@@ -64,11 +58,10 @@ from transformers.utils import is_torchdynamo_compiling, logging
 from transformers import AutoProcessor, BatchFeature
 from qwen_vl_utils.vision_process import smart_resize
 
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.policies.pretrained import PreTrainedPolicy, T
+from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import populate_queues
 from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
+from lerobot.utils.constants import ACTION, OBS_STATE
 
 from lerobot.policies.wall_x.utils import *
 from lerobot.policies.wall_x.constant import *
diff --git a/src/lerobot/policies/wall_x/utils.py b/src/lerobot/policies/wall_x/utils.py
index 10d2ab6c1..19d85aa66 100644
--- a/src/lerobot/policies/wall_x/utils.py
+++ b/src/lerobot/policies/wall_x/utils.py
@@ -21,7 +21,6 @@ Contains data processing utilities, text formatting functions, and helper classe
 for the Wall-X cross-embodiment robotic control model.
 """
 
-import json
 import random
 import re
 from collections import OrderedDict
@@ -626,13 +625,11 @@ def replace_action_token(
         text: List of text strings with action placeholders
         norm_action: Normalized action tensors
         action_tokenizer: Tokenizer for converting actions to tokens
-        dataset_names: Names of datasets for each sample
         dof_masks: Masks for degrees of freedom
 
     Returns:
         List of text strings with action tokens replaced
     """
-    # Filter out multimodal dataset names
     if action_tokenizer is not None and norm_action is not None:
         # Extract actions based on chunk sizes and DOF masks
         norm_action = [