Hi, I really like your work in dealing with multi-agent trajectories prediction. I wen

Time encoder implementation about agentformer HOT 1 OPEN

Dennis-Tsai commented on May 24, 2024

Time encoder implementation

from agentformer.

Comments (1)

PFery4 commented on May 24, 2024

The temporal encoding is managed by the PositionalAgentEncoding class in the agentformer.py file:

AgentFormer/model/agentformer.py

Lines 33 to 99 in e4fe8dd

 class PositionalAgentEncoding(nn.Module): 

 def __init__(self, d_model, dropout=0.1, max_t_len=200, max_a_len=200, concat=False, use_agent_enc=False, agent_enc_learn=False): 

 super(PositionalAgentEncoding, self).__init__() 

 self.dropout = nn.Dropout(p=dropout) 

 self.concat = concat 

 self.d_model = d_model 

 self.use_agent_enc = use_agent_enc 

 if concat: 

 self.fc = nn.Linear((3 if use_agent_enc else 2) * d_model, d_model) 

 pe = self.build_pos_enc(max_t_len) 

 self.register_buffer('pe', pe) 

 if use_agent_enc: 

 if agent_enc_learn: 

 self.ae = nn.Parameter(torch.randn(max_a_len, 1, d_model) * 0.1) 

 else: 

 ae = self.build_pos_enc(max_a_len) 

 self.register_buffer('ae', ae) 

 def build_pos_enc(self, max_len): 

 pe = torch.zeros(max_len, self.d_model) 

 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 

 div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-np.log(10000.0) / self.d_model)) 

 pe[:, 0::2] = torch.sin(position * div_term) 

 pe[:, 1::2] = torch.cos(position * div_term) 

 pe = pe.unsqueeze(0).transpose(0, 1) 

 return pe 

 def build_agent_enc(self, max_len): 

 ae = torch.zeros(max_len, self.d_model) 

 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 

 div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-np.log(10000.0) / self.d_model)) 

 ae[:, 0::2] = torch.sin(position * div_term) 

 ae[:, 1::2] = torch.cos(position * div_term) 

 ae = ae.unsqueeze(0).transpose(0, 1) 

 return ae 

 def get_pos_enc(self, num_t, num_a, t_offset): 

 pe = self.pe[t_offset: num_t + t_offset, :] 

 pe = pe.repeat_interleave(num_a, dim=0) 

 return pe 

 def get_agent_enc(self, num_t, num_a, a_offset, agent_enc_shuffle): 

 if agent_enc_shuffle is None: 

 ae = self.ae[a_offset: num_a + a_offset, :] 

 else: 

 ae = self.ae[agent_enc_shuffle] 

 ae = ae.repeat(num_t, 1, 1) 

 return ae 

 def forward(self, x, num_a, agent_enc_shuffle=None, t_offset=0, a_offset=0): 

 num_t = x.shape[0] // num_a 

 pos_enc = self.get_pos_enc(num_t, num_a, t_offset) 

 if self.use_agent_enc: 

 agent_enc = self.get_agent_enc(num_t, num_a, a_offset, agent_enc_shuffle) 

 if self.concat: 

 feat = [x, pos_enc.repeat(1, x.size(1), 1)] 

 if self.use_agent_enc: 

 feat.append(agent_enc.repeat(1, x.size(1), 1)) 

 x = torch.cat(feat, dim=-1) 

 x = self.fc(x) 

 else: 

 x += pos_enc 

 if self.use_agent_enc: 

 x += agent_enc 

 return self.dropout(x)

Hope this helps!

from agentformer.

Time encoder implementation about agentformer HOT 1 OPEN

Comments (1)

Related Issues (20)

Recommend Projects

React

Vue.js

Typescript

TensorFlow

Django

Laravel

D3

Recommend Topics

javascript

web

server

Machine learning

Visualization

Game

Recommend Org

Facebook

Microsoft

Google

Alibaba

D3

Tencent

	class PositionalAgentEncoding(nn.Module):

	def __init__(self, d_model, dropout=0.1, max_t_len=200, max_a_len=200, concat=False, use_agent_enc=False, agent_enc_learn=False):
	super(PositionalAgentEncoding, self).__init__()
	self.dropout = nn.Dropout(p=dropout)
	self.concat = concat
	self.d_model = d_model
	self.use_agent_enc = use_agent_enc
	if concat:
	self.fc = nn.Linear((3 if use_agent_enc else 2) * d_model, d_model)

	pe = self.build_pos_enc(max_t_len)
	self.register_buffer('pe', pe)
	if use_agent_enc:
	if agent_enc_learn:
	self.ae = nn.Parameter(torch.randn(max_a_len, 1, d_model) * 0.1)
	else:
	ae = self.build_pos_enc(max_a_len)
	self.register_buffer('ae', ae)

	def build_pos_enc(self, max_len):
	pe = torch.zeros(max_len, self.d_model)
	position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
	div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-np.log(10000.0) / self.d_model))
	pe[:, 0::2] = torch.sin(position * div_term)
	pe[:, 1::2] = torch.cos(position * div_term)
	pe = pe.unsqueeze(0).transpose(0, 1)
	return pe

	def build_agent_enc(self, max_len):
	ae = torch.zeros(max_len, self.d_model)
	position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
	div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-np.log(10000.0) / self.d_model))
	ae[:, 0::2] = torch.sin(position * div_term)
	ae[:, 1::2] = torch.cos(position * div_term)
	ae = ae.unsqueeze(0).transpose(0, 1)
	return ae

	def get_pos_enc(self, num_t, num_a, t_offset):
	pe = self.pe[t_offset: num_t + t_offset, :]
	pe = pe.repeat_interleave(num_a, dim=0)
	return pe

	def get_agent_enc(self, num_t, num_a, a_offset, agent_enc_shuffle):
	if agent_enc_shuffle is None:
	ae = self.ae[a_offset: num_a + a_offset, :]
	else:
	ae = self.ae[agent_enc_shuffle]
	ae = ae.repeat(num_t, 1, 1)
	return ae

	def forward(self, x, num_a, agent_enc_shuffle=None, t_offset=0, a_offset=0):
	num_t = x.shape[0] // num_a
	pos_enc = self.get_pos_enc(num_t, num_a, t_offset)
	if self.use_agent_enc:
	agent_enc = self.get_agent_enc(num_t, num_a, a_offset, agent_enc_shuffle)
	if self.concat:
	feat = [x, pos_enc.repeat(1, x.size(1), 1)]
	if self.use_agent_enc:
	feat.append(agent_enc.repeat(1, x.size(1), 1))
	x = torch.cat(feat, dim=-1)
	x = self.fc(x)
	else:
	x += pos_enc
	if self.use_agent_enc:
	x += agent_enc
	return self.dropout(x)