HI, I've implemented the following Autonomous robot navigation system. But in scan callback function I'm trying to update the current location. So, how to send the latest odom_msg to the defined update_current_location(). You can check for any other silly mistakes I made or any possible errors or check for redundancy. Optimize the code. And also give me your suggestions to improve the code and your ideas on what can be done better. Here is the code: #! /usr/bin/env python3 

import tf
import copy
import time
import rospy
import random
import numpy as np
import math
from collections import deque
from sensor_msgs.msg import LaserScan
from ackermann_msgs.msg import AckermannDrive
from nav_msgs.msg import Odometry
import torch
import torch.nn as nn
import torch.optim as optim

FORWARD = 0
LEFT = 1
RIGHT = 2
REVERSE = 3
ACTION_SPACE = [FORWARD, LEFT, RIGHT, REVERSE]

DISTANCE_THRESHOLD = 2  # meters


class DQNAgentPytorch:
    def __init__(self, action_space, observation_space, memory_size=10000, epsilon=1.0, gamma=0.99, learning_rate=0.001):
        self.observation_space = observation_space
        self.action_space = action_space
        self.memory = deque(maxlen=memory_size)
        self.epsilon = epsilon
        self.gamma = gamma
        self.lr = learning_rate
        self.model = self.build_model()
        self.target_model = self.build_model()
        self.optimizer = optim.AdamW(self.model.parameters(), lr=self.lr, amsgrad=True)
        self.criterion = nn.SmoothL1Loss()
        self.update_target_net()

    def build_model(self):
        model = nn.Sequential(
            nn.Linear(self.observation_space.shape[0], 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, len(self.action_space))
        )
        return model

    def add_experience(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def choose_action(self, state, test=False):
        if test or random.random() > self.epsilon:
            with torch.no_grad():
                state_t = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
                q_values = self.model(state_t)
                return torch.argmax(q_values).item()
        else:
            return np.random.choice(self.action_space)

    def train(self, batch_size):
        if len(self.memory) < batch_size:
            return

        minibatch = random.sample(self.memory, batch_size)
        for state, action, reward, next_state, done in minibatch:
            target = self.target_model(torch.tensor(state, dtype=torch.float32).unsqueeze(0)).squeeze(0)
            if done:
                target[action] = reward
            else:
                next_q_values = self.target_model(torch.tensor(next_state, dtype=torch.float32).unsqueeze(0)).squeeze(0)
                target[action] = reward + self.gamma * torch.max(next_q_values)

            inputs = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
            targets = target.unsqueeze(0)

            self.optimizer.zero_grad()

            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()

            # In-place gradient clipping
            torch.nn.utils.clip_grad_value_(self.model.parameters(), 100)
            self.optimizer.step()

    def update_target_net(self):
        self.target_model.load_state_dict(self.model.state_dict())

    def update_epsilon(self, new_epsilon):
        self.epsilon = new_epsilon

    def save_weights(self, file_name):
        torch.save(self.model.state_dict(), file_name)

    def load_weights(self, file_name):
        self.epsilon = 0
        self.model.load_state_dict(torch.load(file_name))

    def calculate_reward(self, state, action, min_distance):
        reward = 0
        if action == REVERSE:
            reward -= 1  # decrease the reward for reversing
        else:
            reward += max(0, (min_distance - DISTANCE_THRESHOLD) / DISTANCE_THRESHOLD)  # increase the reward for avoiding obstacles
        return reward


class EvaderNode:
    def __init__(self, agent):
        rospy.init_node('evader_node', anonymous=True)
        rospy.Subscriber('car_1/scan', LaserScan, self.scan_callback)

        rospy.Subscriber('car_1/base/odom', Odometry, self.odom_callback)
        self.drive_pub = rospy.Publisher('car_1/command', AckermannDrive, queue_size=10)

        # self.current_location = copy.deepcopy(self.init_location)


        self.is_training = True  # Flag to indicate if the agent is being trained


        self.init_location = None  # Set the initial location in odom_callback when available
        self.final_destination = (10, 10)  # Set the final destination here
        self.current_location = None  # Will be updated in update_current_location method
        
        self.distance_threshold = 2  # meters
        self.angle_range = [(-150, -30), (30, 150)]

        self.agent = agent
        self.collision_threshold = 1.0

        self.is_training = True  # Flag to indicate if the agent is being trained
        self.init_odom_received = False

    def process_scan(self, scan):
        # Process the laser scan data to get the state input for the DQN
        state = np.array(scan.ranges)[::10]  # Downsample the scan data for faster processing
        return state

    def check_collision(self, scan):
        for i in range(len(scan.ranges)):
            if scan.ranges[i] < self.collision_threshold:
                return True
        return False
    

    def action_to_drive_msg(self, action):
        drive_msg = AckermannDrive()
        if action == FORWARD:
            drive_msg.speed = 2.0
            drive_msg.steering_angle = 0.0
        elif action == LEFT:
            drive_msg.speed = 2.0
            drive_msg.steering_angle = float(random.uniform(math.pi/4, math.pi/2))
        elif action == RIGHT:
            drive_msg.speed = 2.0
            drive_msg.steering_angle = float(random.uniform(-math.pi/2, -math.pi/4))
        elif action == REVERSE:
            drive_msg.speed = -2.0
            drive_msg.steering_angle = 0.0
        return drive_msg

    # def scan_callback(self, scan):
    #     state = self.process_scan(scan)
    #     action = self.agent.choose_action(state)
    #     drive_msg = self.action_to_drive_msg(action)

    #     min_distance = min(state)
    #     reward = self.agent.calculate_reward(state, action, min_distance)

    #     done = False  # set this flag to True if a termination condition is met
    #     self.agent.add_experience(state, action, reward, state, done)

    #     self.drive_pub.publish(drive_msg)
    
    def reverse_car(self, reverse_distance):
        reverse_drive_msg = AckermannDrive()
        reverse_drive_msg.speed = -2.0
        reverse_drive_msg.steering_angle = 0.0
        reverse_duration = reverse_distance / (-reverse_drive_msg.speed)

        start_time = time.time()
        while time.time() - start_time < reverse_duration:
            self.drive_pub.publish(reverse_drive_msg)
            rospy.sleep(0.1)

    def scan_callback(self, scan):
        state = self.process_scan(scan)
        action = self.agent.choose_action(state, test=not self.is_training)

        collision_detected = self.check_collision(scan)
        if collision_detected:
            self.reverse_car(1.0)  # Set the reverse distance here

        drive_msg = self.action_to_drive_msg(action)
 
        self.update_current_location()

        if not self.is_training:
            if np.linalg.norm(np.array(self.final_destination) - np.array(self.current_location)) < self.distance_threshold:
                self.current_location = copy.deepcopy(self.init_location)
            self.drive_pub.publish(drive_msg)
            return

        min_distance = min(state)
        reward = self.agent.calculate_reward(state, action, min_distance)

        done = False  # set this flag to True if a termination condition is met
        self.agent.add_experience(state, action, reward, state, done)
        self.agent.train(32)  # Set the batch size here

        self.drive_pub.publish(drive_msg)

        # Set the termination condition for training
        if np.linalg.norm(np.array(self.final_destination) - np.array(self.current_location)) < self.distance_threshold:
            self.is_training = False
            self.agent.save_weights('model_weights.pth')
            self.agent.load_weights('model_weights.pth')

    def update_current_location(self, odom_msg):
        position = odom_msg.pose.pose.position
        orientation = odom_msg.pose.pose.orientation
        euler_angles = tf.transformations.euler_from_quaternion((orientation.x, orientation.y, orientation.z, orientation.w))

        self.current_location = (position.x, position.y, euler_angles[2])  # (x, y, yaw)

    def odom_callback(self, odom_msg):

        if not self.init_odom_received:
    
            self.init_location = (odom_msg.pose.pose.position.x,
                                  odom_msg.pose.pose.position.y)
            self.init_odom_received = True

        self.update_current_location(odom_msg)

if __name__ == '__main__':
    try:
        # Initialize the DQNAgentPytorch instance with the necessary environment and parameters
        action_space = ACTION_SPACE
        observation_space = np.zeros(360)  # 360 laser scan equidistant points
        agent = DQNAgentPytorch(action_space, observation_space)
        

        # Initialize the EvaderNode instance with the DQNAgent
        node = EvaderNode(agent)
        
        # Wait for initial odometry information
        while not node.init_odom_received:
            rospy.sleep(0.1)
        
        rospy.spin()

    except rospy.ROSInterruptException:
        pass